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

#include <cstdio>

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

namespace {
std::map<std::string, std::string> modelNames = {
    {"deepseek-chat", "DeepSeek-V3"}, {"deepseek-reasoner", "DeepSeek-R1"}};
}

DeepseekNlpEnginePrivate::DeepseekNlpEnginePrivate()
    : currentModel_("deepseek-chat") {}

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

std::string DeepseekNlpEnginePrivate::engineName() const { return "deepseek"; }

std::string DeepseekNlpEnginePrivate::currentModel() const {
    return currentModel_;
}

std::string DeepseekNlpEnginePrivate::modelInfo() const {
    const char *info = R"(
    {
        "vendor": "深度求索",
        "models": [
            {
                "name": "DeepSeek-V3",
                "model": "deepseek-chat",
                "subConfig": [
                    {
                        "name": "deepseek-chat",
                        "config": [
                            {
                                "friendlyName": "APIKey",
                                "configName": "apiKey"
                            }
                        ]
                    }
                ],
                "maxTokens": 8192,
                "feedbackUrl": "https://platform.deepseek.com/usage"
            },
            {
                "name": "DeepSeek-R1",
                "model": "deepseek-reasoner",
                "subConfig": [
                    {
                        "name": "deepseek-reasoner",
                        "config": [
                            {
                                "friendlyName": "APIKey",
                                "configName": "apiKey"
                            }
                        ]
                    }
                ],
                "maxTokens": 8192,
                "feedbackUrl": "https://platform.deepseek.com/usage"
            }
        ]
    }
    )";

    return info;
}

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

void DeepseekNlpEnginePrivate::setTextGenerationResultCallback(
    ai_engine::gen_ai::nlp::TextGenerationResultCallback callback) {
    textGenerationResultCallback_ = callback;
}

bool DeepseekNlpEnginePrivate::setCurrentModel(const std::string &model) {
    if (modelNames.find(model) == modelNames.end()) {
        return false;
    }

    currentModel_ = model;
    return true;
}
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) {
    DeepseekNlpEnginePrivate *engine = (DeepseekNlpEnginePrivate *)userData;
    auto callback = engine->textGenerationResultCallback_;

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

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

    if (resultJson.isObject() && !resultJson.isNull() &&
        resultJson.isMember("error")) {
        fprintf(stderr, "deepseek 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);
    if (dataStrings[0] == ": keep-alive") {
        return true;
    }
    for (std::string data : dataStrings) {
        data = data.substr(6, data.length() - 1);
        if (data == "[DONE]") {
            break;
        }
        Json::Value obj = deepseek_nlp_util::formatJsonFromString(data);
        Json::Value callbackObj;
        callbackObj["is_end"] =
            static_cast<bool>(obj["choices"][0]["finish_reason"] == "stop");
        if (!obj["choices"][0]["delta"]["reasoning_content"].empty()) {
            callbackObj["reasoning_content"] =
                obj["choices"][0]["delta"]["reasoning_content"];
        }

        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(), modelNames[engine->currentModel_],
            callbackObj.toStyledString()};
        callback(result);
    }
    return true;
}

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

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

void DeepseekNlpEnginePrivate::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 DeepseekNlpEnginePrivate::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;
    }
}

ai_engine::result::AiEngineBaseResult DeepseekNlpEnginePrivate::init() {
    context_["model"] = currentModel_;
    if (getMaxTokens(currentModel_) != -1) {
        context_["max_tokens"] = maxTokens_;
    }
    context_["stream"] = true;
    return {true};
}

ai_engine::result::AiEngineBaseResult DeepseekNlpEnginePrivate::destroy() {
    stopChat();
    return {true};
}

ai_engine::result::AiEngineBaseResult DeepseekNlpEnginePrivate::generateContent(
    const std::string &input) {
    ai_engine::EngineError error{
        (int)ai_engine::NlpEngineErrorCode::FailedToAnswer,
        "Unsupported features"};
    ai_engine::gen_ai::nlp::TextGenerationResult result{
        false,
        error,
        modelNames[currentModel_],
    };
    if (textGenerationResultCallback_) {
        textGenerationResultCallback_(result);
    }

    return {false, error};
}
ai_engine::result::AiEngineBaseResult DeepseekNlpEnginePrivate::chat(
    const std::string &message) {
    tempStreamChatData_.clear();
    isChatRunning_ = true;
    sentenceId_ = 0;

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

    clearContext();

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

    Json::Value context = deepseek_nlp_util::formatJsonFromString(message);
    for (int i = 0; i < context["messages"].size(); i++) {
        addContext_(context["messages"][i]["role"].asString(),
                    context["messages"][i]["content"].asString());
    }
    context_["model"] = currentModel_;
    std::string postString = context_.toStyledString();

    cpr::Url url{"https://api.deepseek.com/chat/completions"};
    cpr::Header header{cpr::Header{{"Content-Type", "application/json"},
                                   {"Authorization", "Bearer " + apiKey_}}};
    cpr::Body body{cpr::Body{{postString}}};

    std::fprintf(stderr, "DeepSeek postString: %s\n", postString.c_str());
    cpr::Response response = cpr::Post(
        url, header, body, cpr::WriteCallback{writeChatData, (intptr_t)this},
        cpr::LowSpeed(20, 10));

    return processChatResponse(response);
}

void DeepseekNlpEnginePrivate::stopGenerateContent() {}

bool DeepseekNlpEnginePrivate::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 DeepseekNlpEnginePrivate::stopChat() {
    if (!isChatRunning_) {
        return;
    }
    if (isStopped_.exchange(true)) {
        return;
    }
    streamChatData_.clear();

    isChatRunning_ = false;
    ai_engine::gen_ai::nlp::TextGenerationResult result{
        true,
        {},
        modelNames[currentModel_],
        std::string{R"({"result":"","is_end":true})"}};
    if (textGenerationResultCallback_ != nullptr) {
        textGenerationResultCallback_(result);
    }
}

ai_engine::result::AiEngineBaseResult
DeepseekNlpEnginePrivate::processChatResponse(cpr::Response &response) {
    isChatRunning_ = false;
    if (isStopped_) {
        isStopped_ = false;
        return {true};
    };

    if (response.error.code == cpr::ErrorCode::REQUEST_CANCELLED) {
        streamChatData_.clear();
        ai_engine::EngineError error(
            (int)deepseek_nlp_server_error::errorCode2nlpResult(
                response.status_code),
            currentErrorInfo_);
        return {false, error};
    }

    if (response.error) {
        streamChatData_.clear();
        fprintf(stderr, "net error: %s\n", response.error.message.c_str());
        ai_engine::EngineError error(
            (int)ai_engine::NlpEngineErrorCode::NetworkDisconnected,
            "HTTP response error");
        return {false, error};
    }

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