/*
 * 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 "baiduvisionengine_p.h"
#include <unistd.h>
#include <map>
#include <cpr/cpr.h>
#include "baiduresponseprocessor.h"
#include "imageloader.h"
#include "servererror.h"
#include "token.h"
#include "util.h"

namespace {
std::map<int, std::string> imageStyleMap = {
    {0, "探索无限"},    {1, "古风"},     {2, "二次元"},      {3, "写实风格"},
    {4, "浮世绘"},      {5, "low poly"}, {6, "未来主义"},    {7, "像素风格"},
    {8, "概念艺术"},    {9, "赛博朋克"}, {10, "洛丽塔风格"}, {11, "巴洛克风格"},
    {12, "超现实主义"}, {13, "水彩画"},  {14, "蒸汽波艺术"}, {15, "油画"},
    {16, "卡通画"},
};
}

BaiduVisionEnginePrivate::BaiduVisionEnginePrivate() {}

BaiduVisionEnginePrivate::~BaiduVisionEnginePrivate() {}

std::string BaiduVisionEnginePrivate::engineName() {
    return std::string("baidu");
}

void BaiduVisionEnginePrivate::setAppId(const std::string &appId) {
    appId_ = appId;
}

void BaiduVisionEnginePrivate::setApiKey(const std::string &apiKey) {
    apiKey_ = apiKey;
}

void BaiduVisionEnginePrivate::setSecretKey(const std::string &secretKey) {
    secretKey_ = secretKey;
}

void BaiduVisionEnginePrivate::setPrompt2ImageCallback(
    ai_engine::lm::vision::ImageResultCallback callback) {
    imageResultCallback_ = callback;
}

void BaiduVisionEnginePrivate::setPrompt2ImageSize(int width, int height) {
    imageWidth_ = width;
    imageHeight_ = height;
}

void BaiduVisionEnginePrivate::setPrompt2ImageNumber(int number) {
    imageNumber_ = number;
}

void BaiduVisionEnginePrivate::setPrompt2ImageStyle(int style) {
    imageStyle_ = imageStyleMap[style];
}

ai_engine::ErrorCode
BaiduVisionEnginePrivate::prompt2Image(const std::string &prompt) {
    std::string taskId = requestPrompt2ImageTask(prompt);
    if (taskId.empty()) {
        return (ai_engine::ErrorCode)resultCode_;
    }
    return processPrompt2ImageTask(taskId);
}

std::string &BaiduVisionEnginePrivate::getAccessToken_() {
    if (accessToken_.empty()) {
        generateAccessToken_();
    }
    return accessToken_;
}

void BaiduVisionEnginePrivate::generateAccessToken_() {
    bool hasNetError = false;
    accessToken_ = token::getBaiduToken(apiKey_, secretKey_, hasNetError,
                                        lastErrorMessage_);
    if (accessToken_.empty()) {
        hasNetError
            ? resultCode_ = (int)ai_engine::ErrorCode::NetError
            : resultCode_ = (int)ai_engine::ErrorCode::AuthenticationFailed;
        fprintf(stderr, "Get baidu vision token failed.\n");
        return;
    }
    resultCode_ = (int)ai_engine::ErrorCode::Success;
}

std::string
BaiduVisionEnginePrivate::requestPrompt2ImageTask(const std::string &prompt) {
    const std::string postUrl{"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/"
                              "txt2imgv2?access_token=" +
                              getAccessToken_()};

    Json::Value postJson;
    if (imageStyle_.empty()) {
        postJson["prompt"] = std::string(prompt);
    } else if (imageStyle_ == "low poly") {
        postJson["prompt"] = std::string(prompt) + "，" + "低多边形";
    } else {
        postJson["prompt"] = std::string(prompt) + "，" + imageStyle_;
    }
    postJson["width"] = imageWidth_;
    postJson["height"] = imageHeight_;
    postJson["image_num"] = imageNumber_;
    std::string postString = postJson.toStyledString();

    cpr::Url url{postUrl};
    cpr::Header header{{"Content-Type", "application/json"},
                       {"Accept", "application/json"}};
    cpr::Body body{postString};
    cpr::Response r = cpr::Post(url, header, body);

    if (r.error) {
        lastErrorMessage_ = r.error.message;
        resultCode_ = (int)ai_engine::ErrorCode::NetError;
        fprintf(stderr, "Net error: %d.\n", (int)r.error.code);
        return std::string();
    }

    Json::Value result = util::formatJsonFromString(r.text);
    if (result.isMember("data") && result["data"].isMember("task_id")) {
        resultCode_ = (int)ai_engine::ErrorCode::Success;
        return result["data"]["task_id"].asString();
    }

    if (token::isBaiduTokenExpiredByResult(result)) {
        generateAccessToken_();
        return requestPrompt2ImageTask(prompt);
    }

    int serverErrorCode = server_error::parseErrorCode(r.text);
    if (serverErrorCode > 0) {
        lastErrorMessage_ = r.text;
        resultCode_ =
            (int)server_error::errorCode2visionResult(serverErrorCode);
        fprintf(stderr, "Get baidu vision result failed %s.\n", r.text.c_str());
    }

    resultCode_ = (int)ai_engine::ErrorCode::RequestFailed;
    return std::string();
}

ai_engine::ErrorCode
BaiduVisionEnginePrivate::processPrompt2ImageTask(const std::string &taskId) {
    while (1) {
        const std::string taskResult = doRequestTaskResult(taskId);
        if (taskResult.empty()) {
            return (ai_engine::ErrorCode)resultCode_;
        }

        ImageDataList list;
        bool isTaskBlocked = false;
        ProcessResult processCode = BaiduResponseProcessor::process(
            taskResult, list, isTaskBlocked, lastErrorMessage_);

        if (processCode == ProcessResult::SUCCEEDED) {
            for (auto &imageData : list) {
                auto imageInfo = vision::image_loader::getImageInfoFromData(
                    imageData.data_.data(), imageData.data_.size());
                std::vector<char> charDataVector;
                for (unsigned char value : imageData.data_) {
                    charDataVector.push_back(static_cast<char>(value));
                }
                ai_engine::lm::vision::ImageData data{
                    imageInfo.format,
                    imageInfo.height,
                    imageInfo.width,
                    (int)list.size(),
                    (int)imageData.index_,
                    charDataVector,
                    imageData.conclusion_ == ApproveConclusion::PASS
                        ? (int)ai_engine::ErrorCode::Success
                        : (int)ai_engine::ErrorCode::ResultBlocked};
                imageResultCallback_(data);
            }
            resultCode_ = (int)ai_engine::ErrorCode::Success;
            return ai_engine::ErrorCode::Success;
        }

        if (processCode == ProcessResult::FAILED) {
            isTaskBlocked
                ? resultCode_ = (int)ai_engine::ErrorCode::ResultBlocked
                : resultCode_ = (int)ai_engine::ErrorCode::RequestFailed;
            return (ai_engine::ErrorCode)resultCode_;
        }
        sleep(1);
    }
}

std::string
BaiduVisionEnginePrivate::doRequestTaskResult(const std::string &taskId) {
    const std::string postUrl{
        "https://aip.baidubce.com/rpc/2.0/ernievilg/v1/getImgv2?access_token=" +
        getAccessToken_()};

    char postData[1024];
    snprintf(postData, sizeof(postData), "{\"task_id\":\"%s\"}",
             taskId.c_str());

    cpr::Url url{postUrl};
    cpr::Header header{{"Content-Type", "application/json"},
                       {"Accept", "application/json"}};
    cpr::Body body{postData};
    cpr::Response r = cpr::Post(url, header, body);

    if (r.error) {
        lastErrorMessage_ = r.error.message;
        resultCode_ = (int)ai_engine::ErrorCode::NetError;
        fprintf(stderr, "Net error: %d.\n", (int)r.error.code);
        return std::string();
    }

    if (token::isBaiduTokenExpiredByResult(r.text)) {
        generateAccessToken_();
        return doRequestTaskResult(taskId);
    }

    int serverErrorCode = server_error::parseErrorCode(r.text);
    if (serverErrorCode > 0) {
        lastErrorMessage_ = r.text;
        resultCode_ =
            (int)server_error::errorCode2visionResult(serverErrorCode);
        fprintf(stderr, "Get baidu vision result failed %s.\n", r.text.c_str());
    }

    resultCode_ = (int)ai_engine::ErrorCode::Success;
    return r.text;
}