/*
 * Copyright 2024 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 "xunfeinlpengine_p.h"
#include <unistd.h>
#include <thread>
#include <curl/websockets.h>

#include "logger.h"
#include "servererror.h"
#include "token.h"
#include "util.h"

bool XunfeiNlpEnginePrivate::connectWebsocket(const std::string &hostUrl) {
    std::string callurl = xunfei_nlp_token::getAuthenticationUrl(apiKey_, secretKey_, hostUrl);

    CURLcode res;
    if (currentCurl_) {
        curl_easy_setopt(currentCurl_, CURLOPT_URL, callurl.c_str());
        curl_easy_setopt(currentCurl_, CURLOPT_CONNECT_ONLY, 2L);
        res = curl_easy_perform(currentCurl_);
        if (res == CURLE_HTTP_RETURNED_ERROR) {
            LOGE("xunfei authenticate failed:", curl_easy_strerror(res));
            currentError_ =
                ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::Unauthorized, "Authentication failed");
            return false;
        }
        if (res != CURLE_OK) {
            LOGE("net error:", res, curl_easy_strerror(res));
            if (res == CURLE_PEER_FAILED_VERIFICATION) {
                currentError_ = ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::SslCertificateError,
                                                       curl_easy_strerror(res));
            } else {
                currentError_ = ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::NetworkDisconnected,
                                                       curl_easy_strerror(res));
            }
            return false;
        }
        return true;
    }
    currentError_ =
        ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::InvalidWords, "Failed to create resource");
    return false;
}

void XunfeiNlpEnginePrivate::closeWebsocket() {
    size_t sent;
    if (currentCurl_) {
        curl_ws_send(currentCurl_, "", 0, &sent, 0, CURLWS_CLOSE);
    }
}

XunfeiNlpEnginePrivate::XunfeiNlpEnginePrivate() {}

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

std::string XunfeiNlpEnginePrivate::engineName() const { return "xunfei"; }

void XunfeiNlpEnginePrivate::setConfig(const std::string &config) {
    Json::Value configJson = xunfei_nlp_util::stringToJson(config);
    if (configJson.isNull() || !configJson.isMember("apiKey") || !configJson.isMember("secretKey") ||
        !configJson.isMember("appId") || !configJson["apiKey"].isString() || !configJson["secretKey"].isString() ||
        !configJson["appId"].isString()) {
        LOGE("Invalid config for xunfei chat engine:", config);
        return;
    }
    appId_ = configJson["appId"].asString();
    apiKey_ = configJson["apiKey"].asString();
    secretKey_ = configJson["secretKey"].asString();

    context_["header"]["app_id"] = appId_;
}

bool XunfeiNlpEnginePrivate::setCurrentModel(const std::string &model) {
    Json::Value modelInfoJson = xunfei_nlp_util::stringToJson(modelInfo());
    for (auto &item : modelInfoJson["models"]) {
        if (item["model"].asString() == model) {
            currentModel_ = model;
            return true;
        }
    }
    LOGE("Invalid current model name:", model);
    return false;
}

std::string XunfeiNlpEnginePrivate::modelInfo() const {
    LOGI("xunfei model info:", "xunfei");
    const char *info = R"(
    {
        "vendor": "讯飞",
        "models": [
            {
                "name": "讯飞-Spark Max",
                "model": "generalv3.5",
                "subConfig": [
                    {
                        "name": "Spark Max",
                        "config": [
                            {
                                "friendlyName": "APPID",
                                "configName": "appId"
                            },
                            {
                                "friendlyName": "APIKey",
                                "configName": "apiKey"
                            },
                            {
                                "friendlyName": "APISecret",
                                "configName": "secretKey"
                            }
                        ]
                    }
                ],
                "maxTokens": 8192,
                "feedbackUrl": "https://console.xfyun.cn/workorder/commit"
            }
        ]
    }
    )";

    return info;
}

void XunfeiNlpEnginePrivate::setTextGenerationResultCallback(
    const ai_engine::gen_ai::nlp::TextGenerationResultCallback &callback) {
    std::lock_guard locker(mutex_);
    textGenerationResultCallback_ = callback;
}

bool XunfeiNlpEnginePrivate::doSendData(const char *data, int dataLength, int type) {
    if (!isChatRunning()) {
        LOGI("XunfeiNlpEnginePrivate::chat: no chat is running, stop to send "
             "data");
        return false;
    }
    size_t sent;
    CURLcode resultCode = curl_ws_send(currentCurl_, data, dataLength, &sent, 0, type);
    if (resultCode != CURLE_OK) {
        LOGE("send data to xunfei failed:", curl_easy_strerror(resultCode));
        currentError_ = ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::NetworkDisconnected,
                                               curl_easy_strerror(resultCode));
        return false;
    }
    return true;
}

CURLcode XunfeiNlpEnginePrivate::testConnectToServer() {
    std::string callurl = xunfei_nlp_token::getAuthenticationUrl(apiKey_, secretKey_, hostUrl_);
    CURL *curl = curl_easy_init();
    curl_easy_setopt(curl, CURLOPT_URL, callurl.c_str());

    curl_easy_setopt(curl, CURLOPT_CONNECT_ONLY, 2L); /* websocket style */

    /* Perform the request, res gets the return code */
    CURLcode res = curl_easy_perform(curl);
    curl_easy_cleanup(curl);

    return res;
}

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

void XunfeiNlpEnginePrivate::clearContext() {
    if (context_["payload"].isMember("message")) {
        context_["payload"]["message"]["text"].clear();
    }
}

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

ai_engine::result::AiEngineBaseResult XunfeiNlpEnginePrivate::buildChatMessage(const std::string &message) {
    Json::Value context = xunfei_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_["parameter"]["chat"]["top_k"] = context["top_k"].asInt();
    }

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

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

ai_engine::result::AiEngineBaseResult XunfeiNlpEnginePrivate::init() {
    currentError_ = ai_engine::EngineError();
    context_["parameter"]["chat"]["domain"] = "generalv3.5";
    if (getMaxTokens(currentModel_) != -1) {
        context_["parameter"]["chat"]["max_tokens"] = maxTokens_;
    }
    return {true, ai_engine::EngineError()};
}

ai_engine::result::AiEngineBaseResult XunfeiNlpEnginePrivate::generateContent(const std::string &input) {
    Json::Value inputJson = xunfei_nlp_util::stringToJson(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 XunfeiNlpEnginePrivate::chat(const std::string &message) {
    currentError_ = ai_engine::EngineError();

    {
        std::lock_guard locker(mutex_);
        if (textGenerationResultCallback_ == nullptr) {
            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()) {
        LOGI("XunfeiNlpEnginePrivate::chat: no chat is running");
        return {true, currentError_};
    }

    if (!initResources() || !connectWebsocket(hostUrl_)) {
        if (isStopped_) {
            return {true, ai_engine::EngineError()};
        }
        decreaseChatCounter();
        return {false, currentError_};
    }

    ai_engine::result::AiEngineBaseResult result = buildChatMessage(message);
    if (!result.success) {
        return result;
    }
    std::string postString = xunfei_nlp_util::jsonToString(context_);
    LOGI(postString);
    if (!doSendData(postString.c_str(), postString.size(), CURLWS_TEXT)) {
        removeLastContext();
        streamChatData_.clear();
        closeWebsocket();
        if (isStopped_) {
            return {true, ai_engine::EngineError()};
        }
        decreaseChatCounter();
        return {false, currentError_};
    }

    handleReceiveData();

    streamChatData_.clear();

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

bool XunfeiNlpEnginePrivate::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 XunfeiNlpEnginePrivate::stopChat() {
    if (!isChatRunning()) {
        LOGI("XunfeiNlpEnginePrivate::stopChat: chat is not running");
        return;
    }
    streamChatData_.clear();

    decreaseChatCounter();

    isStopped_ = true;
    ai_engine::gen_ai::nlp::TextGenerationResult result{true, currentError_, currentModelName_,
                                                        std::string{R"({"result":"","is_end":true})"}};
    {
        std::lock_guard locker(mutex_);
        if (textGenerationResultCallback_ != nullptr) {
            textGenerationResultCallback_(result);
        }
    }
    closeWebsocket();
    releaseResources();
}

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

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

void XunfeiNlpEnginePrivate::threadSafeCallChatResultCallback(
    const ai_engine::gen_ai::nlp::TextGenerationResult &result) const {
    std::lock_guard locker(mutex_);
    if (textGenerationResultCallback_ != nullptr) {
        textGenerationResultCallback_(result);
    }
}

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

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

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

void XunfeiNlpEnginePrivate::handleReceiveData() {
    int retryCount = 0;
    while (!isStopped_) {
        int retCode;
        const std::string chatData = receiveChatData(&retCode);

        if (isNetworkError(retCode)) {
            handleNetworkError(retCode);
            break;
        }
        if (needReReceiveData(retCode)) {
            ++retryCount;
            if (retryCount > 1000 * 10) {
                auto ret = testConnectToServer();
                if (ret != 0) {
                    handleNetworkError(ret);
                    break;
                } else {
                    retryCount = 0;
                }
            }
            usleep(1000);
            continue;
        }
        if (int errorCode = xunfei_nlp_server_error::parseErrorCode(chatData); errorCode != 0) {
            handleServerError(errorCode, chatData);
            break;
        }

        if (isChatDataReady(retCode)) {
            bool isFinished = false;
            handleChatData(chatData, &isFinished);
            if (isFinished) {
                decreaseChatCounter();
                break;
            }
        }
        retryCount = 0;
        usleep(1000);
    }
    if (isStopped_) {
        isStopped_ = false;
    }
}
std::string XunfeiNlpEnginePrivate::receiveChatData(int *retCode) const {
    size_t rlen;
    const curl_ws_frame *meta;
    char buffer[4100]{0};
    CURLcode curlCode = curl_ws_recv(currentCurl_, buffer, sizeof(buffer), &rlen, &meta);
    *retCode = curlCode;
    return std::string(buffer, rlen);
}

bool XunfeiNlpEnginePrivate::isNetworkError(int retCode) const {
    return retCode != CURLE_OK && retCode != CURLE_AGAIN && retCode != CURLE_RECV_ERROR;
}

bool XunfeiNlpEnginePrivate::isChatDataReady(int retCode) const { return retCode == CURLE_OK; }

bool XunfeiNlpEnginePrivate::needReReceiveData(int retCode) const {
    return retCode == CURLE_AGAIN || retCode == CURLE_RECV_ERROR;
}

void XunfeiNlpEnginePrivate::handleNetworkError(int retCode) {
    if (isStopped_) {
        isStopped_ = false;
        return;
    }
    LOGE("Net error:", curl_easy_strerror(CURLcode(retCode)));
    currentError_ = ai_engine::EngineError((int)ai_engine::NlpEngineErrorCode::NetworkDisconnected,
                                           curl_easy_strerror(CURLcode(retCode)));
    ai_engine::gen_ai::nlp::TextGenerationResult result{false, currentError_, currentModelName_, std::string{}};
    threadSafeCallChatResultCallback(result);
}

void XunfeiNlpEnginePrivate::handleServerError(int errorCode, const std::string &errorMessage) {
    if (isStopped_) {
        isStopped_ = false;
        return;
    }
    LOGE("Xunfei chat failed", errorMessage);
    currentError_ = ai_engine::EngineError((int)xunfei_nlp_server_error::errorCode2nlpResult(errorCode), errorMessage);
    ai_engine::gen_ai::nlp::TextGenerationResult result{false, currentError_, currentModelName_, std::string{}};
    threadSafeCallChatResultCallback(result);
}

void XunfeiNlpEnginePrivate::handleChatData(const std::string &chatData, bool *isFinished) {
    Json::Value obj = xunfei_nlp_util::stringToJson(chatData);

    Json::Value callbackObj;
    callbackObj["sentence_id"] = obj["payload"]["choices"]["seq"];
    callbackObj["is_end"] = (obj["header"]["status"].asInt() == 2);
    callbackObj["result"] = obj["payload"]["choices"]["text"][0]["content"];
    callbackObj["finish_reason"] = "";
    streamChatData_.push_back(obj["payload"]["choices"]["text"][0]["content"]);
    ai_engine::gen_ai::nlp::TextGenerationResult result{
        true,
        ai_engine::EngineError(),
        currentModelName_,
        callbackObj.toStyledString(),
    };
    threadSafeCallChatResultCallback(result);
    *isFinished = (obj["header"]["status"].asInt() == 2);
}

bool XunfeiNlpEnginePrivate::initResources() {
    if (currentCurl_) {
        curl_easy_cleanup(currentCurl_);
        currentCurl_ = nullptr;
    }
    currentCurl_ = curl_easy_init();

    return true;
}

void XunfeiNlpEnginePrivate::releaseResources() {
    if (currentCurl_) {
        curl_easy_cleanup(currentCurl_);
        currentCurl_ = nullptr;
    }
}
