/*
 * Copyright (C) 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 "embeddingserver.h"

#include <string>

#include "embeddingproxy.h"
#include "task.h"

namespace ai_engine::core_ai::embedding {

EmbeddingServer::EmbeddingServer()
    : embeddingProxy_(std::make_unique<EmbeddingProxy>()) {
    startServer_ = true;
    processTaskThread_ =
        std::thread(&EmbeddingServer::processTaskInThread, this);
}
EmbeddingServer::~EmbeddingServer() {
    startServer_ = false;
    taskQueueCv_.notify_all();
    if (processTaskThread_.joinable()) processTaskThread_.join();
}
void EmbeddingServer::setEmbeddingResultCallback(
    EmbeddingResultCallback callback) {
    callback_ = std::move(callback);
}
void EmbeddingServer::processTaskInThread() {
    while (startServer_) {
        EmbeddingTask task;
        std::unique_lock<std::mutex> lock(mutex_);

        taskQueueCv_.wait(
            lock, [this]() { return !tasks_.empty() || !startServer_; });

        if (!startServer_) {
            std::cout << "[Embedding server]: Embedding servcer stop"
                      << std::endl;
            break;
        }

        std::unique_lock<std::mutex> lockTask(taskMutex_);
        if (!tasks_.empty()) {
            task = tasks_.top();
            tasks_.pop();
        }

        // 阻塞到有任务之后，开始进行向量化任务
        EmbeddingProxyResult result = doEmbeddingInThread(task);
        if (callbacks_.find(task.id()) != callbacks_.end()) {
            std::cout << "[Embedding server]:Task callback id: " << task.id()
                      << "found " << std::endl;
            callbacks_[task.id()](EmbeddingResult{
                result.success,
                ai_engine::EngineError(result.error.code, result.error.message),
                "ondevice embedding engine", result.res});
            callbacks_.erase(task.id());
        }
    }
}

EmbeddingProxyResult EmbeddingServer::doEmbeddingInThread(
    const EmbeddingTask &task) {
    EmbeddingProxyResult result;
    switch (task.type()) {
        case EmbeddingTask::EmbeddingDataType::Text:
            std::cout << "[Embedding server]:Embedding text start" << std::endl;
            result = embeddingProxy_->embeddingText(task.getData());
            break;
        case EmbeddingTask::EmbeddingDataType::TextByImageModel:
            std::cout
                << "[Embedding server]:Embedding text with imageModel start"
                << std::endl;
            result =
                embeddingProxy_->embeddingTextWithImageModel(task.getData());
            break;
        case EmbeddingTask::EmbeddingDataType::Image:
            std::cout
                << "[Embedding server]:Embedding image from imageFile start"
                << std::endl;
            result = embeddingProxy_->embeddingImage(task.getData(), 0);
            break;
        case EmbeddingTask::EmbeddingDataType::ImageBase64:
            std::cout << "[Embedding server]:Embedding image from "
                         "base6ImageData start"
                      << std::endl;
            result = embeddingProxy_->embeddingImage(task.getData(), 1);
            break;
        default:
            std::cerr << "[Embedding server]:Task type error!" << std::endl;
    }
    return result;
}

bool EmbeddingServer::addTask(const EmbeddingTask &task) {
    if (task.type() == EmbeddingTask::EmbeddingDataType::Text) {
        // 文本向量化模型任务
        if (!initTextEmbeddingModule().success) {
            std::cerr << "[Embedding server]:Text embedding module init error!"
                      << std::endl;
            return false;
        }

    } else {
        // 图像向量化模型任务
        if (!initImageEmbeddingModule().success) {
            std::cerr << "[Embedding server]:Image embedding module init error!"
                      << std::endl;
            return false;
        }
    }

    std::unique_lock<std::mutex> lock(taskMutex_);
    callbacks_.insert(make_pair(task.id(), callback_));

    tasks_.push(task);

    taskQueueCv_.notify_all();
    return true;
}
std::string EmbeddingServer::modelInfo() const {
    return embeddingProxy_->modelInfo();
}

bool EmbeddingServer::createTritonClient() {
    return embeddingProxy_->createTritonClient();
}
result::AiEngineBaseResult EmbeddingServer::initTextEmbeddingModule() {
    if (!createTritonClient()) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(CREATE_TRITON_CLIENT_ERROR,
                                          "Create triton client error!")};
    }

    EmbeddingProxyResult result = embeddingProxy_->loadTextModel();
    return result::AiEngineBaseResult{
        result.success,
        ai_engine::EngineError(result.error.code, result.error.message)};
}
result::AiEngineBaseResult EmbeddingServer::destroyTextEmbeddingModule() {
    // 目前不需要我们自己卸载模型，所以这里是写死的
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}
result::AiEngineBaseResult EmbeddingServer::embeddingText(
    const std::string &text) {
    if (!callback_) {
        std::cerr << "[Embedding server]:Callback is nullptr!" << std::endl;
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Callback is nullptr!")};
    }

    bool res =
        addTask(EmbeddingTask(text, EmbeddingTask::EmbeddingDataType::Text));
    if (!res) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Add embedding task error!")};
    }
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}

result::AiEngineBaseResult EmbeddingServer::initImageEmbeddingModule() {
    if (!createTritonClient()) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(CREATE_TRITON_CLIENT_ERROR,
                                          "Create triton client error!")};
    }

    EmbeddingProxyResult result = embeddingProxy_->loadImageModel();
    return result::AiEngineBaseResult{
        result.success,
        ai_engine::EngineError(result.error.code, result.error.message)};
}
result::AiEngineBaseResult EmbeddingServer::destroyImageEmbeddingModule() {
    // 目前不需要我们自己卸载模型，所以这里是写死的
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}
result::AiEngineBaseResult EmbeddingServer::embeddingImageFromImageFile(
    const std::string &filePath) {
    if (!callback_) {
        std::cerr << "[Embedding server]:Callback is nullptr!" << std::endl;
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Callback is nullptr!")};
    }
    bool res = addTask(
        EmbeddingTask(filePath, EmbeddingTask::EmbeddingDataType::Image));
    if (!res) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Add embedding task error!")};
    }
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}
result::AiEngineBaseResult EmbeddingServer::embeddingImageFromBase64ImageData(
    const std::string &fileData) {
    if (!callback_) {
        std::cerr << "[Embedding server]:Callback is nullptr!" << std::endl;
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Callback is nullptr!")};
    }
    bool res = addTask(
        EmbeddingTask(fileData, EmbeddingTask::EmbeddingDataType::ImageBase64));
    if (!res) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Add embedding task error!")};
    }
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}
result::AiEngineBaseResult EmbeddingServer::embeddingTextWithImageModel(
    const std::string &text) {
    if (!callback_) {
        std::cerr << "[Embedding server]:Callback is nullptr!" << std::endl;
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Callback is nullptr!")};
    }
    bool res = addTask(EmbeddingTask(
        text, EmbeddingTask::EmbeddingDataType::TextByImageModel));
    if (!res) {
        return result::AiEngineBaseResult{
            false, ai_engine::EngineError(ADD_EMBEDDING_TASK_ERROR,
                                          "Add embedding task error!")};
    }
    return result::AiEngineBaseResult{true,
                                      ai_engine::EngineError(0, "Success")};
}

}  // namespace ai_engine::core_ai::embedding