#include "agentrequesthandler.h"
#include "toolparser.h"
#include "networkdefs.h"
#include "mcpserver.h"

#include <QJsonDocument>
#include <QJsonObject>
#include <QRegularExpression>
#include <QDebug>

#include <QLoggingCategory>
Q_DECLARE_LOGGING_CATEGORY(logAgent)

namespace uos_ai {

AgentRequestHandler::AgentRequestHandler(QObject *parent)
    : Agent(parent)
    , m_mcpClient(new McpClient(this))
{
    
}

AgentRequestHandler::~AgentRequestHandler()
{
    
}

QSharedPointer<MCPServer> AgentRequestHandler::mcpServer() const
{
    return QSharedPointer<MCPServer>(new MCPServer(m_name));
}

bool AgentRequestHandler::initClient()
{
    if (!m_mcpClient->init()) {
        qCWarning(logAgent) << "Failed to initialize MCP client";
        return false;
    }

    return true;
}

QStringList AgentRequestHandler::listServers() const
{
    return syncCall<QStringList>([this]() {
        qCDebug(logAgent) << "list servers in" << QThread::currentThreadId();
        return m_mcpClient->listServers(m_name);
    });
}

QPair<int, QJsonValue> AgentRequestHandler::listTools() const
{
    return syncCall<QPair<int, QJsonValue>>([this]() {
        qCDebug(logAgent) << "list tools in" << QThread::currentThreadId();
        return m_mcpClient->getTools(m_name);
    });
}

QPair<int, QString> AgentRequestHandler::fetchTools(const QStringList &servers)
{
    m_tools = QJsonArray();
    qCDebug(logAgent) << "Fetching tools from servers:" << servers;

    QPair<int, QJsonValue> srv = syncCall<QPair<int, QJsonValue>>([this, servers]() {
        qCDebug(logAgent) << "query servers in" << QThread::currentThreadId();
        return m_mcpClient->queryServers(m_name, servers);
    });

    if (srv.first != 0)
        return qMakePair(srv.first, QString());
    
    QJsonArray toolsArray;
    QString error;
    QStringList invaildSrv = servers;
    for (const QJsonValue &serverValue : srv.second.toArray()) {
        if (!serverValue.isObject()) {
            continue;
        }

        QJsonObject serverObj = serverValue.toObject();
        QString name = serverObj.value("name").toString();

        if (name.isEmpty() || (!servers.isEmpty() && !servers.contains(name)))
            continue;

        invaildSrv.removeOne(name);
        if (serverObj.contains("error")) {
            QString er = serverObj.value("error").toString();
            error.append(QString("%0: %1\n\n").arg(name).arg(er));
            continue;
        }

        if (serverObj.contains("tools") && serverObj["tools"].isArray()) {
            QJsonArray serverTools = serverObj["tools"].toArray();
            for (const QJsonValue &toolValue : serverTools) {
                toolsArray.append(toolValue);
            }
        }
    }

    for (const QString &name : invaildSrv)
         error.append(QString("%0: %1\n\n").arg(name).arg("No such server."));

    if (!error.isEmpty()) {
        error = tr("MCP server is not available") + QString("\n\n") + error;
        return qMakePair(AIServer::ErrorType::MCPSeverUnavailable, error);
    }

    m_tools = std::move(toolsArray);
    return qMakePair(AIServer::ErrorType::NoError, QString());
}


void AgentRequestHandler::processRequest(const QJsonObject &question, const QJsonArray &history, QSharedPointer<LLM> llm, const QVariantHash &params)
{
    // 先刷新一次服务
    {
        auto ret = syncCall<QPair<int, QJsonObject>>([this]() {
            qCDebug(logAgent) << "sync servers in" << QThread::currentThreadId();
            return m_mcpClient->syncServers(m_name);
        });

        if (ret.first == 0) {
            auto details = ret.second.value("details").toObject();
            if (!details.isEmpty())
                qCInfo(logAgent) << "Synced servers" << details;
        } else {
            qCWarning(logAgent) << "Failed to sync severs:" << ret.second;
        }
    }

    // 取消
    if (canceled) {
        qCDebug(logAgent) << "agent canceled before fetch tools.";
        return;
    }

    {
        auto toolRet = fetchTools(params.value(PREDICT_PARAM_MCPSERVERS).toStringList());
        if (toolRet.first != 0) {
            qCWarning(logAgent) << "Failed to fetch tools:" << toolRet.second;
            llm->setLastError(toolRet.first);
            llm->setLastErrorString(toolRet.second);
            return;
        }
    }

    // 取消
    if (canceled){
        qCDebug(logAgent) << "agent canceled before init  messages.";
        return;
    }

    auto messages = initChatMessages(question, history);
    llm->switchStream(true);

    while (!canceled) {
        QString current_output;
        QString tool_results;
        ToolParser parser;

        connect(llm.data(), &LLM::readyReadChatDeltaContent, this, [llm, &current_output, &tool_results, &parser, this](const QString &content) {
            if (canceled) {
                qCDebug(logAgent()) << "received delta content after abort." << content;
                return;
            }

            bool ret = handleModelOutput(content, current_output, tool_results, parser);
            if (!ret) {
                canceled = true;
                // cancle
                llm->aborted();
            }
        });

        llm->predict(QJsonDocument(messages).toJson(), {});

        if (!canceled) {
            parser.finalize();
            for (auto action : parser.getResults()) {
                if (action.first == ToolParser::Text) {
                    if (!action.second.isEmpty())
                        this->textChainContent(action.second);
                } else if (action.first == ToolParser::Function) {
                    qCritical(logAgent) << "Catch function in end:" << action.second;
                }
            }
        }

        if (tool_results.isEmpty()) {
            qCInfo(logAgent) << "No tool results, end chat. messages size:" << messages.size();
            break;
        }

        messages.append(QJsonObject{{"role","assistant"}, {"content", current_output}});
        messages.append(QJsonObject{{"role", "user"}, {"content",tool_results}});

        disconnect(llm.data(), &LLM::readyReadChatDeltaContent, this, nullptr);
    }
}

void AgentRequestHandler::cancel()
{
    canceled = true;
}

void AgentRequestHandler::textChainContent(const QString &content)
{
    QJsonObject message;

    message.insert("content", content);
    message.insert("chatType", ChatAction::ChatTextPlain);  // 普通文本类型

    QJsonObject wrapper;
    wrapper.insert("message", message);
    wrapper.insert("stream", true);

    emit readyReadChatDeltaContent(QJsonDocument(wrapper).toJson());
}

void AgentRequestHandler::toolUseContent(const ToolUse &tool)
{
    QJsonObject message = tool.toJson();
    message.insert("chatType", ChatAction::ChatToolUse);

    QJsonObject wrapper;
    wrapper.insert("message", message);
    wrapper.insert("stream", true);

    emit readyReadChatDeltaContent(QJsonDocument(wrapper).toJson());
}

bool AgentRequestHandler::handleModelOutput(const QString &content, QString &current_output, QString &tool_results, ToolParser &parser)
{
    QString realContent;
    QJsonDocument doc = QJsonDocument::fromJson(content.toUtf8());
     if (!doc.isNull() && doc.isObject()) {
         QJsonObject obj = doc.object();
         if (obj.contains("message") && obj["message"].isObject()) {
             QJsonObject messageObj = obj["message"].toObject();
             if (messageObj.value("chatType").toInt() != ChatAction::ChatTextPlain) {
                 emit readyReadChatDeltaContent(content);
                 return true;
             }

             if (messageObj.contains("content")) {
                 realContent = messageObj["content"].toString();
             }
         }
     }
     if (realContent.isEmpty())
        return true;

    current_output += realContent;
    parser.feed(realContent);

    for (auto action : parser.getResults()) {
        if (action.first == ToolParser::Text) {
            if (!action.second.isEmpty())
                this->textChainContent(action.second);
        } else if (action.first == ToolParser::Function) {
            QString func_content = action.second;
            if (!func_content.isEmpty()) {
                QString toolName;
                QJsonObject args;
                if (ToolParser::toJson(func_content, toolName, args)) {
                    if (m_toolList.contains(toolName)) {
                        qCInfo(logAgent) << "Agent" << name() << "calling tool:" << func_content << QThread::currentThreadId();

                        ToolUse tool;
                        tool.content = ToolParser::restoreFunction(func_content);
                        tool.name = toolName;
                        tool.params = QString::fromUtf8(QJsonDocument(args).toJson());
                        tool.index = m_usedTool.size();
                        this->toolUseContent(tool);

        //                QString expl = args["explanation"].toString();
        //                args.remove("explanation");

                        auto result = syncCall<QPair<int, QString>>([this, toolName, args]() {
                            qCDebug(logAgent) << "call tool in" << QThread::currentThreadId();
                            return m_mcpClient->callTool(m_name, toolName, args);
                        });

                        tool.result = result.second;
                        bool isAbort = (result.first != AIServer::ErrorType::NoError)
                                     && (result.first != AIServer::ErrorType::MCPToolError);
                        tool.status = !result.first ? ToolUse::Completed : ToolUse::Failed;
                        m_usedTool.append(tool);
                        this->toolUseContent(tool);

                        if (!isAbort) {
                            if (result.first) {
                                qCWarning(logAgent) << "The tool" << toolName
                                                    << " was successfully called, but the tool execution was incorrect"
                                                    << result.second;
                            } else {
                                qCInfo(logAgent) << "Tool" << toolName << "executed successfully";
                            }

                            tool_results += QString("<tool_output>%0</tool_output>\n").arg(result.second);
                        } else {
                            qCWarning(logAgent) << "Tool execution failed for" << toolName << result.second;
                            return false;
                        }
                    } else {
                        qCWarning(logAgent) << "Invalid tool requested:" << toolName;
                    }
                } else {
                    qCWarning(logAgent) << "Failed to parse function content:" << func_content;
                }
            }
        }
    }

    return true;
}

} // namespace uos_ai 
