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

#include <jsoncpp/json/json.h>

#include "modelconfig/modelconfigmanager.h"
#include "modelconfig/modelstatustracker.h"
#include "taskmanager/taskrunner.h"
#include "utils/logger.h"
#include "utils/util.h"

ImageGenerator::ImageGenerator(std::unique_ptr<ImageGenerationEngine> visionEngine)
    : visionEngine_(std::move(visionEngine)) {
    initVisionEngineResultCallback();
}

ImageGenerator::~ImageGenerator() {
    if (visionEngine_ != nullptr) {
        visionEngine_->setGenerationResultCallback(nullptr);
    }
    // TODO 析构不执行，后续查找原因
    LOGD("~ImageGenerator");
}

void ImageGenerator::setImageGeneratorCallback(const ImageGeneratorResultCallback &callback) {
    imageGeneratorResultCallback_ = callback;
}

void ImageGenerator::setVisionEngine(std::unique_ptr<ImageGenerationEngine> visionEngine) {
    visionEngine_ = std::move(visionEngine);
    initVisionEngineResultCallback();
    setVisionEngineConfig(imageConfig_);
}

void ImageGenerator::generateImage(const std::string &prompt, ImageGeneratorResultCallback resultCallback,
                                   ImageGenerationFinishedCallback finisedCallback) {
    ImageGenerationRequest request{prompt, resultCallback, finisedCallback};
    requestQueue_.enqueue(request);
    generateImageAsync();
}

void ImageGenerator::stopImageGeneration() {
    visionEngine_->stopImageGeneration();
    engineInUse_ = false;
}

int ImageGenerator::getSupportedImageNumber() { return visionEngine_->supportedImageNumber(); }

std::string ImageGenerator::getSupportedImageSize() {
    std::vector<std::pair<int, int>> sizes = visionEngine_->supportedImageSizes();
    Json::Value value;
    for (const auto &[width, height] : sizes) {
        value.append(Json::Value(std::to_string(width) + "x" + std::to_string(height)));
    }
    return util::jsonToString(value);
}

std::string ImageGenerator::getSupportedImageStyle() {
    std::vector<std::string> styles = visionEngine_->supportedImageStyles();
    Json::Value value;
    for (const std::string &style : styles) {
        value.append(style);
    }
    return util::jsonToString(value);
}

void ImageGenerator::setImageConfig(const genai::ImageConfig &config) {
    imageConfig_ = config;
    LOGI("setImageConfig {} {} {} {} {}", imageConfig_.width, imageConfig_.height, imageConfig_.number,
         imageConfig_.size, imageConfig_.style);

    setVisionEngineConfig(config);
}

void ImageGenerator::initVisionEngineResultCallback() {
    if (visionEngine_ != nullptr) {
        visionEngine_->setGenerationResultCallback(
            [this](const ai_engine::gen_ai::vision::ImageGenerationResult &imageData) {
                onVisionEngineResultCallback(imageData);
            });
    }
}

void ImageGenerator::onVisionEngineResultCallback(const ai_engine::gen_ai::vision::ImageGenerationResult &imageData) {
    LOGD("Call prompt2ImageCallback");
    ImageGenerationResult result{};
    result.format = genai::getVisionImageFormat(imageData.format);
    result.height = imageData.height;
    result.width = imageData.width;
    result.total = imageData.total;
    result.index = imageData.index;
    result.imageBase64Data = imageData.imageBase64Data;
    result.errorCode = genai::getVisionErrorCode(imageData.error.getCode());
    result.errorMessage = imageData.error.getMessage();
    if (currentRequest_.resultCallback != nullptr) {
        currentRequest_.resultCallback(result);
    }
}

void ImageGenerator::setVisionEngineConfig(const genai::ImageConfig &config) {
    visionEngine_->setImageNumber(config.number);
    visionEngine_->setImageSize(config.width, config.height);
    visionEngine_->setImageStyle(config.style);
}

void ImageGenerator::onErrorOccurred(const ai_engine::EngineError &error) {
    ImageGenerationResult result{};
    result.format = int(genai::ImageFormat::IMAGE_FORMAT_UNKNOWN);
    result.width = 0;
    result.height = 0;
    result.total = 0;
    result.index = -1;
    result.errorCode = genai::getVisionErrorCode(error.getCode());
    result.errorMessage = error.getMessage();
    if (currentRequest_.resultCallback) {
        currentRequest_.resultCallback(result);
    }
}

int ImageGenerator::getImageNumberFromJsonString(const std::string &params) {
    Json::Value root;
    Json::Reader reader;

    bool parsingSuccessful = reader.parse(params, root);
    if (!parsingSuccessful) {
        LOGE("Failed to parse JSON: {}", reader.getFormattedErrorMessages());
        return 0;
    }

    if (!root.isMember("image-number")) {
        LOGE("Key 'image-number' not found in JSON");
        return 0;
    }
    int imageNumber = root["image-number"].asInt();
    return imageNumber;
}

void ImageGenerator::generateImageAsync() {
    if (taskId_ != task::INVALID_TASK_ID) {
        return;
    }
    taskId_ = task::generateTaskId();
    std::shared_ptr sharedObject = getSharedObject();
    task::TaskRunner::getInstance().addTask({taskId_, task::TaskType::TextToImage, [sharedObject]() {
                                                 sharedObject->generateImageTask();
                                                 sharedObject->taskId_ = task::INVALID_TASK_ID;
                                             }});
}

void ImageGenerator::generateImageTask() {
    while (requestQueue_.wait_dequeue_timed(currentRequest_, 10 * 1000 * 1000)) {
        ModelStatusTracker statusTracker(visionEngine_->currentModelName());
        engineInUse_ = true;
        auto result = visionEngine_->generateImage(currentRequest_.prompt);
        if (!result.success) {
            LOGE("Generate image error info: {}, {}.", result.error.getCode(), result.error.getMessage());
            onErrorOccurred(result.error);
        }
        engineInUse_ = false;
        if (currentRequest_.finishedCallback) {
            currentRequest_.finishedCallback();
        }
    }
}
