#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;
}
