#include "eappaiprompt.h"
#include "eparserdocument.h"
#include "embeddingserver.h"
#include "global.h"

#include <QJsonArray>
#include <QStandardPaths>

EConversationPrompt::EConversationPrompt(const QString &userData)
    : EAiPrompt(userData)
{
}

QString EConversationPrompt::getAiPrompt()
{
    QString output;

    switch (m_llm) {
    case CHATGPT_3_5:
    case CHATGPT_3_5_16K:
    case CHATGPT_4:
    case CHATGPT_4_32K:
//        output += "You are playing the role of an intelligent AI. I will input a piece of content, and your task is to give a best reply based on the input content and provide the output in the same language as the input.";
//        output += "\nInput:\n";
//        output +=  m_userParam + "\n";
//        output += "Output:";
        output += m_userParam;
        break;
    default:
        output += m_userParam;
        break;
    }

    return output;
}

EAiDocSummaryPrompt::EAiDocSummaryPrompt(const QString &userData, const QString &aiData)
    : EAiPrompt(userData, aiData)
{
}

QString EAiDocSummaryPrompt::getAiPrompt()
{
    QString promptTemplate = "你是一个文档助手，根据文档内容回答问题。\n" \
                     "文档内容：%1 \n" \
                     "问题：%2 \n";
    QString aiPrompt = promptTemplate.arg(m_aiParam).arg(m_userParam);

    return lengthValid(aiPrompt);
}

EAiSysAssistantPrompt::EAiSysAssistantPrompt(const QString &userData)
    : EAiPrompt(userData)
{
    embeddingSearch();
}

QString EAiSysAssistantPrompt::getAiPrompt()
{
    QString promptTemplate = QCoreApplication::translate("EAiPrompt", "---Role---\n" \
                                                  "You are a helpful assistant, answering questions about the background knowledge provided.\n\n" \
                                                  "---Goal---\n" \
                                                  "Respond to users' questions, incorporating any relevant common sense.If you don't know the answer, just say so. Don't make it up.\n\n" \
                                                  "---knowledge---\n" \
                                                  "%1\n\n" \
                                                  "---question---\n" \
                                                  "%2\n");
    QString context = promptTemplate.arg(m_aiParam).arg(m_userParam);

    return context;
}

QString EAiSysAssistantPrompt::getReferences()
{
    QJsonObject reference;
    reference.insert("type", uos_ai::ExtentionType::KnowledgeBase);

    QJsonArray docArray;
    for (const QString &doc : m_references.keys()) {
        QJsonObject docObj;
        docObj.insert("docPath", doc);

        QJsonArray docContents;
        for (const QString &docContent : m_references.value(doc)) {
            docContents << docContent;
        }
        docObj.insert("docContents", docContents);

        docArray.append(docObj);
    }
    reference.insert("sources", docArray);

    return QJsonDocument(reference).toJson(QJsonDocument::Compact);
}

void EAiSysAssistantPrompt::embeddingSearch()
{
    QString resultData = EmbeddingServer::getInstance().embeddingSearch(m_userParam, 5, AssistantType::UOS_SYSTEM_ASSISTANT /* uos / deepin */);
    QJsonObject resultObj = QJsonDocument::fromJson(resultData.toUtf8()).object();

    QString knowleadge;
    for (auto res : resultObj["result"].toArray()) {
        QString docPath = res.toObject().value("source").toString();
        QString docContent = res.toObject().value("content").toString();

        m_references[docPath].push_back(docContent);
        knowleadge += docContent;
    }
    m_aiParam = knowleadge;
}

EAiPersonalAssistantPrompt::EAiPersonalAssistantPrompt(const QString &userData)
    : EAiPrompt(userData)
{
    embeddingSearch();
}

QString EAiPersonalAssistantPrompt::getAiPrompt()
{
    QString context;
    if (m_references.isEmpty()) {
        QString promptTemplate = QCoreApplication::translate("EAiPrompt", "The output answer starts with \"No relevant information was found in your knowledge base.\"" \
                                                      "---question--\n" \
                                                      ":%1\n");
        context = promptTemplate.arg(m_userParam);
    } else {
        QString promptTemplate = QCoreApplication::translate("EAiPrompt", "---Role---\n" \
                                                      "You are a helpful assistant, answering questions about the background knowledge provided.\n\n" \
                                                      "---Goal---\n" \
                                                      "Respond to users' questions, incorporating any relevant common sense.If you don't know the answer, just say so. Don't make it up.\n\n" \
                                                      "---knowledge---\n" \
                                                      "%1\n\n" \
                                                      "---question---\n" \
                                                      "%2");
        context = promptTemplate.arg(m_aiParam).arg(m_userParam);
    }

    return context;
}

QString EAiPersonalAssistantPrompt::getReferences()
{
    QJsonObject reference;
    reference.insert("type", uos_ai::ExtentionType::KnowledgeBase);

    QJsonArray docArray;
    for (const QString &doc : m_references.keys()) {
        QJsonObject docObj;
        docObj.insert("docPath", doc);

        QJsonArray docContents;
        for (const QString &docContent : m_references.value(doc)) {
            docContents << docContent;
        }
        docObj.insert("docContents", docContents);

        docArray.append(docObj);
    }
    reference.insert("sources", docArray);

    return QJsonDocument(reference).toJson(QJsonDocument::Compact);
}

void EAiPersonalAssistantPrompt::embeddingSearch()
{
    QString resultData = EmbeddingServer::getInstance().embeddingSearch(m_userParam, 5, AssistantType::PERSONAL_KNOWLEDGE_ASSISTANT);
    QJsonObject resultObj = QJsonDocument::fromJson(resultData.toUtf8()).object();

    QString knowleadge;
    QHash<QString, QVector<QString>> docs;
    for (auto res : resultObj["result"].toArray()) {
        QString docPath = res.toObject().value("source").toString();
        QString docContent = res.toObject().value("content").toString();

        m_references[docPath].push_back(docContent);
        knowleadge += docContent;
    }
    m_aiParam = knowleadge;
}

EAiGenerateQuestionPrompt::EAiGenerateQuestionPrompt()
{
}

QString EAiGenerateQuestionPrompt::getAiPrompt()
{
    QStringList docContents = EmbeddingServer::getInstance().getDocContent();
    QString chunk;
    int it = 1;
    for (const QString &docContent : docContents) {
        chunk += QString::number(it) + ":" + docContent + "\n";
        it++;
    }

    QString context;
    QString promptTemplate ="你是问答任务的助手。\n" \
                            "使用以下背景知识生成问题，中文一组，英文一组，两种语言各20个问题。输出格式为JSON，示例如下：\n" \
                            "```json\n{\n    \"chinese\": [\n    \"问题1\",\n    \"问题2\"\n]\n\"english\": [\n    \"question1\",\n    \"question2\"\n]\n}```\n" \
                            "背景知识：%1\n" \
                            "问题：";
    context = promptTemplate.arg(chunk);

    return context;
}

EAiWordSelectionPrompt::EAiWordSelectionPrompt(const QString &prompt)
    : EAiPrompt(prompt)
{
}

QString EAiWordSelectionPrompt::getAiPrompt()
{
    return m_userParam;
}

EAiInstructionPrompt::EAiInstructionPrompt(const QString &prompt)
    : EAiPrompt(prompt)
{
}

QString EAiInstructionPrompt::getAiPrompt()
{
    return m_userParam;
}

EAiWritingPrompt::EAiWritingPrompt(const QString &userData)
    : EAiPrompt(userData)
{
}

QString EAiWritingPrompt::getAiPrompt()
{
    QString promptTemplate = QCoreApplication::translate("EAiPrompt", "---Role---\n" \
                                                  "You are a professional writer with expertise in various writing styles and formats.\n\n" \
                                                  "---Goal---\n" \
                                                  "Based on the user's request, create high-quality content that meets their specific needs.\n\n" \
                                                  "---Writing Types---\n" \
                                                  "1. Article: Write well-structured articles with clear arguments and engaging content\n" \
                                                  "2. Speeches: Create compelling speeches with proper opening, body, and conclusion\n\n" \
                                                  "3. Outlines: Create a structured outline for the given topic\n" \
                                                  "4. Notifications: Write formal notices or announcements\n" \
                                                  "5. Posts: Create engaging social media content\n" \
                                                  "6. Work Report: Write comprehensive work reports\n" \
                                                  "7. Research Report: Create detailed research reports\n" \
                                                  "---Input---\n" \
                                                  "%1\n\n" \
                                                  "---Output Format---\n" \
                                                  "Please provide the content in the following format:\n" \
                                                  "Title:\n" \
                                                  "Content:");

    return promptTemplate.arg(m_userParam);
}

EAiTextProcessingPrompt::EAiTextProcessingPrompt(const QString &userData)
    : EAiPrompt(userData)
{
}

QString EAiTextProcessingPrompt::getAiPrompt()
{
    QString promptTemplate = QCoreApplication::translate("EAiPrompt", "---Role---\n" \
                                                  "You are a text processing expert with deep expertise in language analysis, writing techniques, and text optimization.\n\n" \
                                                  "---Goal---\n" \
                                                  "Process and improve the input text while maintaining its original meaning and style.\n\n" \
                                                  "---Processing Types---\n" \
                                                  "1. Summary: Create a concise summary of the main points\n" \
                                                  "2. Proofread: Identify and fix grammatical, spelling, and stylistic errors\n" \
                                                  "3. Explain: Provide detailed explanations of complex concepts or passages\n" \
                                                  "4. Expand: Elaborate on key points while maintaining the original style\n" \
                                                  "5. Continue: Continue the text in a coherent and natural way\n" \
                                                  "6. Polish: Enhance the text's clarity, flow, and impact while preserving its essence\n\n" \
                                                  "---Input---\n" \
                                                  "%1\n\n" \
                                                  "---Output Format---\n" \
                                                  "Please provide the processed text directly without any additional characters or formatting.");

    return promptTemplate.arg(m_userParam);
}

EAiTranslationPrompt::EAiTranslationPrompt(const QString &userData)
    : EAiPrompt(userData)
{
}

QString EAiTranslationPrompt::getAiPrompt()
{
    QString promptTemplate = QCoreApplication::translate("EAiPrompt", "---Role---\n" \
                                                  "You are a professional translator. Your task is to translate the given text accurately and naturally.\n\n" \
                                                  "---Goal---\n" \
                                                  "Translate the input text while maintaining its original meaning, tone, and style. If the input is in Chinese, translate it to English; if it's in English, translate it to Chinese.\n\n" \
                                                  "---Input---\n" \
                                                  "%1\n\n" \
                                                  "---Output Format---\n" \
                                                  "Please provide the translation directly without any additional characters or formatting.");

    return promptTemplate.arg(m_userParam);
}
