/*
 * Copyright 2025 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "qwennlpengine_p.h"

#include <cstdio>

#include "servererror.h"
#include "util.h"

QwenNlpEnginePrivate::QwenNlpEnginePrivate() {}

QwenNlpEnginePrivate::~QwenNlpEnginePrivate() { streamChatData_.clear(); }

std::string QwenNlpEnginePrivate::engineName() const { return "qwen"; }

std::string QwenNlpEnginePrivate::modelInfo() const {
    const char *info = R"(
    {
        "vendor": "通义千问",
        "models": [
            {
                "name": "通义千问-qwen-plus",
                "model": "qwen-plus",
                "subConfig": [
                    {
                        "name": "qwen-plus",
                        "config": [
                            {
                                "friendlyName": "APIKey",
                                "configName": "apiKey"
                            }
                        ]
                    }
                ],
                "maxTokens": 8192,
                "feedbackUrl": "https://www.aliyun.com/service"
            }
        ]
    }
    )";

    return info;
}

void QwenNlpEnginePrivate::setConfig(const std::string &config) {
    Json::Value configJson = qwen_nlp_util::formatJsonFromString(config);
    if (configJson.isNull() || !configJson.isMember("apiKey") || !configJson["apiKey"].isString()) {
        std::fprintf(stderr, "Invalid config for qwen nlp engine: %s\n", config.c_str());
        return;
    }
    apiKey_ = configJson["apiKey"].asString();
}

bool QwenNlpEnginePrivate::setCurrentModel(const std::string &model) {
    const Json::Value modelInfoJson = qwen_nlp_util::formatJsonFromString(modelInfo());
    for (const auto &item : modelInfoJson["models"]) {
        if (item["model"].asString() == model) {
            currentModel_ = model;
            return true;
        }
    }
    fprintf(stderr, "Invalid current model name: %s\n", model.c_str());
    return false;
}

void QwenNlpEnginePrivate::setTextGenerationResultCallback(
    const ai_engine::gen_ai::nlp::TextGenerationResultCallback &callback) {
    textGenerationResultCallback_ = callback;
}

namespace {
std::vector<std::string> splitDataString(const std::string &dataStr) {
    std::vector<std::string> datas;

    int oldIndex = 0;
    int newIndex = 0;
    while ((newIndex = dataStr.find_first_of("\n\n", newIndex)) != std::string::npos) {
        datas.push_back(dataStr.substr(oldIndex, newIndex - oldIndex));
        newIndex += 2;
        oldIndex = newIndex;
    }
    return datas;
}
}  // namespace

bool writeChatData(std::string data, intptr_t userData) {
    QwenNlpEnginePrivate *engine = (QwenNlpEnginePrivate *)userData;
    ai_engine::gen_ai::nlp::TextGenerationResultCallback callback = engine->textGenerationResultCallback_;

    if (engine->isStopped_ || !engine->isChatRunning()) {
        return false;
    }

    Json::Value resultJson = qwen_nlp_util::formatJsonFromString(data);

    if (resultJson.isObject() && !resultJson.isNull() && resultJson.isMember("error")) {
        fprintf(stderr, "qwen chat failed: %s\n", data.c_str());
        engine->currentErrorInfo_ = data;
        return false;
    }

    std::string resultString = engine->tempStreamChatData_.data() + data;
    if (resultString.substr(resultString.length() - 2) != "\n\n") {
        engine->tempStreamChatData_ += data;
        return true;
    }
    engine->tempStreamChatData_.clear();

    std::vector<std::string> dataStrings = splitDataString(resultString);
    for (std::string data : dataStrings) {
        data = data.substr(6, data.length() - 1);
        if (data == "[DONE]") {
            break;
        }
        Json::Value obj = qwen_nlp_util::formatJsonFromString(data);
        Json::Value callbackObj;
        callbackObj["is_end"] = static_cast<bool>(obj["choices"][0]["finish_reason"] == "stop");
        callbackObj["result"] = obj["choices"][0]["delta"]["content"];
        callbackObj["sentence_id"] = engine->sentenceId_;
        engine->sentenceId_ += 1;
        engine->streamChatData_.push_back(callbackObj["result"]);
        ai_engine::gen_ai::nlp::TextGenerationResult result{true, ai_engine::EngineError(), engine->currentModelName_,
                                                            callbackObj.toStyledString()};
        callback(result);
    }
    return true;
}

int QwenNlpEnginePrivate::getMaxTokens(const std::string &model) {
    Json::Value modelInfoJson = qwen_nlp_util::formatJsonFromString(modelInfo());
    for (auto &item : modelInfoJson["models"]) {
        if (item["model"].asString() == model) {
            maxTokens_ = item["maxTokens"].asInt();
            return maxTokens_;
        }
    }
    return -1;
}

void QwenNlpEnginePrivate::clearContext() {
    if (context_.isMember("messages")) {
        context_["messages"].clear();
    }
}

void QwenNlpEnginePrivate::addContext(const std::string &role, const std::string &content) {
    if (role.empty() || content.empty()) {
        return;
    }
    if (role == "system") {
        context_["messages"][0]["role"] = role;
        context_["messages"][0]["content"] = content;
    } else {
        uint32_t size = context_["messages"].size();
        context_["messages"][size]["role"] = role;
        context_["messages"][size]["content"] = content;
    }
}

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::buildChatMessage(const std::string &message) {
    Json::Value context = qwen_nlp_util::stringToJson(message);
    ai_engine::result::AiEngineBaseResult result;
    if (!context.isMember("messages")) {
        result.success = false;
        result.error = {int(ai_engine::NlpEngineErrorCode::InvalidParamType),
                        "The input information lacks messages content."};
        return result;
    }
    if (!context["messages"].isArray()) {
        result.success = false;
        result.error = {int(ai_engine::NlpEngineErrorCode::InvalidParamType), "The messages is not array."};
        return result;
    }
    uint32_t size = context["messages"].size();
    for (uint32_t i = 0; i < size; ++i) {
        Json::Value item = context["messages"][i];
        if ((!item.isMember("role")) || (!item["role"].isString())) {
            continue;
        }
        if ((!item.isMember("content")) || (!item["content"].isString())) {
            continue;
        }
        std::string role = item["role"].asString();
        std::string content = item["content"].asString();
        addContext(role, content);
    }
    if (context.isMember("top_k") && context["top_k"].isInt()) {
        context_["top_k"] = context["top_k"].asInt();
    }

    result.success = true;
    result.error = ai_engine::EngineError{};
    return result;
}

void QwenNlpEnginePrivate::removeLastContext() {
    int size = context_["messages"].size();
    if (size - 1 >= 0 && context_["messages"][size - 1]["role"].asString() == "user") {
        Json::Value removeValue;
        context_["messages"].removeIndex(size - 1, &removeValue);
        return;
    }
    if ((size - 1) >= 0 && (size - 2) >= 0 && context_["messages"][size - 1]["role"].asString() == "assistant" &&
        context_["messages"][size - 2]["role"].asString() == "user") {
        Json::Value removeValue;
        context_["messages"].removeIndex(size - 1, &removeValue);
        context_["messages"].removeIndex(size - 2, &removeValue);
        return;
    }
}

void QwenNlpEnginePrivate::decreaseChatCounter() {
    if (chatCounter_ > 0) {
        --chatCounter_;
    }
}

void QwenNlpEnginePrivate::increaseChatCounter() { ++chatCounter_; }

bool QwenNlpEnginePrivate::isChatRunning() const { return chatCounter_ > 0; }

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::init() {
    currentError_ = ai_engine::EngineError();

    context_["model"] = currentModel_;
    context_["stream"] = true;
    if (getMaxTokens(currentModel_) != -1) {
        context_["max_tokens"] = maxTokens_;
    }
    context_["enable_search"] = true;

    return {true, ai_engine::EngineError()};
}

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::generateContent(const std::string &input) {
    Json::Value inputJson = qwen_nlp_util::formatJsonFromString(input);
    if (input.empty() || inputJson.empty() || !inputJson.isMember("prompt")) {
        currentError_ =
            ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::TooLittleData, "The message is empty");
        return {false, currentError_};
    }
    std::string promptContent = inputJson["prompt"].asString();
    Json::Value messageJson;
    messageJson["messages"] = Json::Value(Json::arrayValue);
    Json::Value userMessage;
    userMessage["role"] = "user";
    userMessage["content"] = promptContent;
    messageJson["messages"].append(userMessage);

    std::string message = messageJson.toStyledString();
    return chat(message);
}

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::chat(const std::string &message) {
    tempStreamChatData_.clear();
    currentError_ = ai_engine::EngineError();
    sentenceId_ = 0;

    if (!textGenerationResultCallback_) {
        return {false, currentError_};
    }
    if (message.empty()) {
        currentError_ =
            ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::TooLittleData, "The message is empty");
        return {false, currentError_};
    }

    clearContext();
    increaseChatCounter();

    if (!waitForStopped()) {
        isStopped_ = false;
    }

    if (!isChatRunning()) {
        printf("QwenNlpEnginePrivate::chat: no chat is running");
        return {true, currentError_};
    }

    ai_engine::result::AiEngineBaseResult result = buildChatMessage(message);
    if (!result.success) {
        return result;
    }
    std::string postString = context_.toStyledString();
    printf("post string %s", postString.c_str());

    cpr::Url url{"https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"};
    cpr::Header header{cpr::Header{{"Content-Type", "application/json"}, {"Authorization", "Bearer " + apiKey_}}};
    cpr::Body body{cpr::Body{{postString}}};
    cpr::Response response =
        cpr::Post(url, header, body, cpr::WriteCallback{writeChatData, (intptr_t)this}, cpr::LowSpeed(20, 10));

    return processChatResponse(response);
}

bool QwenNlpEnginePrivate::waitForStopped() {
    auto startTime = std::chrono::steady_clock::now();
    while (isStopped_) {
        if (std::chrono::steady_clock::now() - startTime > std::chrono::seconds(10)) {
            return false;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }
    return true;
}

void QwenNlpEnginePrivate::stopChat() {
    if (!isChatRunning()) {
        printf("QwenNlpEnginePrivate::stopChat: chat is not running");
        return;
    }
    decreaseChatCounter();

    isStopped_ = true;
    streamChatData_.clear();

    ai_engine::gen_ai::nlp::TextGenerationResult result{true, currentError_, currentModelName_,
                                                        std::string{R"({"result":"","is_end":true})"}};

    if (textGenerationResultCallback_ != nullptr) {
        textGenerationResultCallback_(result);
    }
}

void QwenNlpEnginePrivate::stopGenerateContent() { isStopped_ = true; }

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::destroy() {
    stopChat();
    return {true, ai_engine::EngineError()};
}

ai_engine::result::AiEngineBaseResult QwenNlpEnginePrivate::processChatResponse(cpr::Response &response) {
    if (isStopped_ || !isChatRunning()) {
        isStopped_ = false;
        return {true, ai_engine::EngineError()};
    }
    decreaseChatCounter();
    if (response.error.code == cpr::ErrorCode::REQUEST_CANCELLED) {
        streamChatData_.clear();
        currentError_ = ai_engine::EngineError(
            (int)qwen_nlp_server_error::errorCode2nlpResult(response.status_code, currentErrorInfo_),
            currentErrorInfo_);
        return {false, currentError_};
    }

    if (response.error) {
        streamChatData_.clear();
        fprintf(stderr, "net error: %s\n", response.error.message.c_str());
        if (response.error.code == cpr::ErrorCode::SSL_REMOTE_CERTIFICATE_ERROR) {
            currentError_ =
                ai_engine::EngineError(int(ai_engine::NlpEngineErrorCode::SslCertificateError), response.error.message);
        } else {
            currentError_ =
                ai_engine::EngineError(int(ai_engine::NlpEngineErrorCode::NetworkDisconnected), response.error.message);
        }
        return {false, currentError_};
    }

    streamChatData_.clear();
    return {true, currentError_};
}