/*
 * 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 "customnlpengine_p.h"

#include <cstdio>

#include "servererror.h"
#include "util.h"
namespace {
std::vector<std::string> splitDataString(const std::string &dataStr,
                                         const std::string &endStr) {
    std::vector<std::string> datas;

    int oldIndex = 0;
    int newIndex = 0;
    while ((newIndex = dataStr.find_first_of(endStr, newIndex)) !=
           std::string::npos) {
        datas.push_back(dataStr.substr(oldIndex, newIndex - oldIndex));
        newIndex += endStr.length();
        oldIndex = newIndex;
    }
    return datas;
}
}  // namespace
bool writeChatData(std::string data, intptr_t userData) {
    CustomNlpEnginePrivate *engine = (CustomNlpEnginePrivate *)userData;
    ai_engine::gen_ai::nlp::TextGenerationResultCallback callback =
        engine->textGenerationResultCallback_;

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

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

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

    if (data == ": keep-alive\n\n") {
        return true;
    }

    std::string resultString = engine->tempStreamChatData_.data() + data;

    if (resultString.length() <= 5) {
        engine->tempStreamChatData_ += data;
        return true;
    }

    if (!engine->isValidResponse(resultString)) {
        engine->currentErrorInfo_ =
            "Unrecognized response format, parsing failed, response is \n" +
            resultString;
        return false;
    }

    if (data.compare(0, 7, "data: {") == 0) {
        engine->endStr_ = "\n\n";
    }
    if (resultString.substr(resultString.length() - engine->endStr_.length()) !=
        engine->endStr_) {
        engine->tempStreamChatData_ += data;
        return true;
    }
    engine->tempStreamChatData_.clear();
    std::vector<std::string> dataStrings =
        splitDataString(resultString, engine->endStr_);

    for (std::string data : dataStrings) {
        if (engine->endStr_ == "\n\n") {
            data = data.substr(6, data.length() - 1);
        }
        if (data == "[DONE]") {
            break;
        }
        Json::Value callbackObj = engine->formatResultJsonFromInputString(data);
        if (callbackObj.isNull()) {
            continue;
        }
        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;
}

CustomNlpEnginePrivate::CustomNlpEnginePrivate() {}

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

std::string CustomNlpEnginePrivate::engineName() const { return "custom"; }

std::string CustomNlpEnginePrivate::modelInfo() const {
    const char *info = R"(
    {
        "vendor": "自定义模型",
        "models": [
            {
            }
        ]
    }
    )";

    return info;
}

void CustomNlpEnginePrivate::setConfig(const std::string &config) {
    Json::Value configJson = custom_nlp_util::formatJsonFromString(config);
    if (configJson.isNull() || !configJson.isMember("apiKey") ||
        !configJson["apiKey"].isString() || !configJson.isMember("apiUrl") ||
        !configJson["apiUrl"].isString() || !configJson.isMember("modelName") ||
        !configJson["modelName"].isString() ||
        !configJson.isMember("modelVersion") ||
        !configJson["modelVersion"].isString()) {
        std::fprintf(stderr, "Invalid config for custom nlp engine: %s\n",
                     config.c_str());
        return;
    }
    apiKey_ = configJson["apiKey"].asString();
    buildChatAndContentGenerationUrl(configJson["apiUrl"].asString());
    currentModel_ = configJson["modelVersion"].asString();
    currentModelName_ = configJson["modelName"].asString();
}

void CustomNlpEnginePrivate::buildChatAndContentGenerationUrl(
    const std::string &apiUrl) {
    const std::string chatStr = "chat/completions";
    const std::string generateContentStr = "completions";
    generateContentUrl_ = apiUrl;
    chatUrl_ = apiUrl;
    if (apiUrl.length() >= generateContentStr.length() &&
        apiUrl.substr(apiUrl.length() - generateContentStr.length()) !=
            generateContentStr) {
        if (apiUrl.back() == '/') {
            generateContentUrl_ = apiUrl + generateContentStr;
        } else {
            generateContentUrl_ = apiUrl + "/" + generateContentStr;
        }
    }
    if (apiUrl.length() >= chatStr.length() &&
        apiUrl.substr(apiUrl.length() - chatStr.length()) != chatStr) {
        if (apiUrl.back() == '/') {
            chatUrl_ = apiUrl + chatStr;
        } else {
            chatUrl_ = apiUrl + "/" + chatStr;
        }
    }
}

bool CustomNlpEnginePrivate::setCurrentModel(const std::string &model) {
    return true;
}

void CustomNlpEnginePrivate::setTextGenerationResultCallback(
    ai_engine::gen_ai::nlp::TextGenerationResultCallback callback) {
    textGenerationResultCallback_ = std::move(callback);
}

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

void CustomNlpEnginePrivate::addContext_(const std::string &role,
                                         const std::string &message) {
    if (role == "system") {
        if (!message.empty()) {
            context_["messages"][0]["role"] = role;
            context_["messages"][0]["content"] = message;
        }
        return;
    }
    int size = context_["messages"].size();
    context_["messages"][size]["role"] = role;
    context_["messages"][size]["content"] = message;
}

void CustomNlpEnginePrivate::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 CustomNlpEnginePrivate::decreaseChatCounter() {
    if (chatCounter_ > 0) {
        --chatCounter_;
    }
}

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

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

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

    context_["model"] = currentModel_;
    context_["stream"] = true;

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

ai_engine::result::AiEngineBaseResult CustomNlpEnginePrivate::generateContent(
    const std::string &input) {
    tempStreamChatData_.clear();
    isGenerateContentCalled_ = true;
    isThinking_ = false;
    hasReasoningContent_ = false;
    sentenceId_ = 0;
    currentError_ = ai_engine::EngineError();

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

    clearContext();

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

    Json::Value context = custom_nlp_util::formatJsonFromString(input);
    context_["prompt"] = context["prompt"].asString();
    std::string postString = context_.toStyledString();

    cpr::Url url{generateContentUrl_};
    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);
}

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

    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("CustomNlpEnginePrivate::chat: no chat is running");
        return {true, currentError_};
    }

    Json::Value context = custom_nlp_util::formatJsonFromString(message);
    for (int i = 0; i < context["messages"].size(); i++) {
        addContext_(context["messages"][i]["role"].asString(),
                    context["messages"][i]["content"].asString());
    }

    if (context.isMember("top_k") && context["top_k"].isInt()) {
        context_["top_k"] = context["top_k"].asInt();
    }

    std::string postString = context_.toStyledString();

    printf("postString: %s", postString.c_str());

    cpr::Url url{chatUrl_};
    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::ConnectTimeout{10000},  // 连接超时10秒
        cpr::LowSpeed{20, 10});

    return processChatResponse_(response);
}

bool CustomNlpEnginePrivate::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 CustomNlpEnginePrivate::stopChat() {
    if (!isChatRunning()) {
        printf("CustomNlpEnginePrivate::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 CustomNlpEnginePrivate::stopGenerateContent() {
    if (isStopped_.exchange(true)) {
        return;
    }
    if (isGenerateContentCalled_) {
        isGenerateContentCalled_ = false;
        ai_engine::gen_ai::nlp::TextGenerationResult result{
            true, currentError_, currentModelName_,
            std::string{R"({"result":"","is_end":true})"}};

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

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

bool CustomNlpEnginePrivate::isValidResponse(
    const std::string &response) const {
    if (response.length() >= 9 && response.compare(0, 5, "data:") != 0 &&
        response.compare(0, 6, "\ndata:") != 0 &&
        response.compare(0, 9, R"({"model":)") != 0) {
        fprintf(stderr, "Invalid response: %s\n", response.c_str());
        return false;
    }
    return true;
}

Json::Value CustomNlpEnginePrivate::formatResultJsonFromInputString(
    const std::string &inputString) {
    Json::Value inputResultJson =
        custom_nlp_util::formatJsonFromString(inputString);

    if (preProcessResultJson(inputResultJson)) {
        return Json::Value();
    }

    Json::Value resultJson;
    resultJson["is_end"] = isResultEnd(inputResultJson);

    if (isThinking()) {
        resultJson["reasoning_content"] = getThinkingResult(inputResultJson);
    } else {
        resultJson["result"] = getResult(inputResultJson);
    }
    resultJson["sentence_id"] = sentenceId_;
    sentenceId_ += 1;

    return resultJson;
}

bool CustomNlpEnginePrivate::preProcessResultJson(
    const Json::Value &inputResultJson) {
    const std::string startThinkingStr = "\u003cthink\u003e";
    const std::string endThinkingStr = "\u003c/think\u003e";
    if (isOpenAiApiResult(inputResultJson)) {
        if (inputResultJson["choices"][0]["delta"].isMember(
                "reasoning_content") &&
            !inputResultJson["choices"][0]["delta"]["reasoning_content"]
                 .empty()) {
            isThinking_ = true;
            hasReasoningContent_ = true;
            return false;
        }
        if (inputResultJson["choices"][0]["delta"].isMember("reasoning") &&
            !inputResultJson["choices"][0]["delta"]["reasoning"].empty()) {
            isThinking_ = true;
            hasReasoningContent_ = true;
            return false;
        }
        if (hasReasoningContent_) {
            hasReasoningContent_ = false;
            isThinking_ = false;
        }

        if (inputResultJson["choices"][0]["delta"]["content"].asString() ==
                startThinkingStr ||
            inputResultJson["choices"][0]["text"].asString() ==
                startThinkingStr) {
            isThinking_ = true;
            return true;
        }

        if (inputResultJson["choices"][0]["delta"]["content"].asString() ==
                endThinkingStr ||
            inputResultJson["choices"][0]["text"].asString() ==
                endThinkingStr) {
            isThinking_ = false;
            return true;
        }
    }

    if (isOllamaApiResult(inputResultJson)) {
        if (inputResultJson["message"]["content"].asString() ==
            startThinkingStr) {
            isThinking_ = true;
            return true;
        }

        if (inputResultJson["message"]["content"].asString() ==
            endThinkingStr) {
            isThinking_ = false;
            return true;
        }
    }

    return false;
}

bool CustomNlpEnginePrivate::isOpenAiApiResult(
    const Json::Value &result) const {
    return result.isObject() && result.isMember("choices");
}

bool CustomNlpEnginePrivate::isOllamaApiResult(
    const Json::Value &result) const {
    return result.isObject() && result.isMember("message");
}

bool CustomNlpEnginePrivate::isResultEnd(const Json::Value &result) const {
    if (isOpenAiApiResult(result)) {
        return result["choices"][0]["finish_reason"] == "stop" ||
               result["choices"][0]["finish_reason"] == "length";
    }

    if (isOllamaApiResult(result)) {
        return result["done"].asBool();
    }

    return false;
}

bool CustomNlpEnginePrivate::isThinking() const { return isThinking_; }

std::string CustomNlpEnginePrivate::getThinkingResult(
    const Json::Value &result) const {
    if (isOpenAiApiResult(result)) {
        if (!result["choices"][0]["delta"]["reasoning_content"].empty()) {
            return result["choices"][0]["delta"]["reasoning_content"]
                .asString();
        }
        if (!result["choices"][0]["delta"]["reasoning"].empty()) {
            return result["choices"][0]["delta"]["reasoning"].asString();
        }
        if (!result["choices"][0]["text"].empty()) {
            return result["choices"][0]["text"].asString();
        }
        return result["choices"][0]["delta"]["content"].asString();
    }

    if (isOllamaApiResult(result)) {
        return result["message"]["content"].asString();
    }

    return "";
}

std::string CustomNlpEnginePrivate::getResult(const Json::Value &result) const {
    if (isOpenAiApiResult(result)) {
        return result["choices"][0].isMember("delta")
                   ? result["choices"][0]["delta"]["content"].asString()
               : result["choices"][0].isMember("text")
                   ? result["choices"][0]["text"].asString()
                   : "";
    }

    if (isOllamaApiResult(result)) {
        return result["message"]["content"].asString();
    }

    return "";
}
ai_engine::result::AiEngineBaseResult
CustomNlpEnginePrivate::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();
        auto errorCode =
            custom_nlp_server_error::errorCode2nlpResult(response.status_code);
        currentError_ = ai_engine::EngineError(static_cast<int>(errorCode),
                                               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),
                "HTTP response error");
        }
        return {false, currentError_};
    }
    if (response.status_code < 200 || response.status_code >= 300) {
        streamChatData_.clear();
        auto errorCode =
            custom_nlp_server_error::errorCode2nlpResult(response.status_code);
        currentError_ = ai_engine::EngineError(static_cast<int>(errorCode),
                                               "HTTP status error");
        return {false, currentError_};
    }

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