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

using namespace std;
#ifdef TRITON_CLIENT
namespace tc = triton::client;
#endif

MattingTritonClient::MattingTritonClient(const string &url) {
#ifdef TRITON_CLIENT
    // Create the Triton client once during object construction
    FAIL_IF_ERR(tc::InferenceServerHttpClient::Create(&client_, url, false),
                "unable to create http client");
#endif
}

cv::Mat MattingTritonClient::preprocess(cv::Mat &m_srcImg) {
    cv::Mat dstimg;
    cv::cvtColor(m_srcImg, dstimg, cv::COLOR_BGR2RGB);
    dstimg.convertTo(dstimg, CV_32FC3, 1.0f / 255.0f);
    cv::Mat mean_mat(dstimg.size(), CV_32FC3,
                     cv::Scalar(this->mean[0], this->mean[1], this->mean[2]));
    cv::subtract(dstimg, mean_mat, dstimg);
    cv::Mat std_mat(dstimg.size(), CV_32FC3,
                    cv::Scalar(this->std[0], this->std[1], this->std[2]));
    cv::divide(dstimg, std_mat, dstimg);
    cv::resize(dstimg, dstimg, cv::Size(inputWidth, inputHeight),
               cv::INTER_LINEAR);
    dstimg = cv::dnn::blobFromImage(dstimg);
    return dstimg;
}

ImgBase64Result MattingTritonClient::MattingRequestByPath(string ImagePath) {
    ImgBase64Result res;
    // 使用 ifstream 检查文件是否存在
    std::ifstream file(ImagePath, std::ios::binary);
    if (!file.is_open()) {
        // 如果文件无法打开，返回错误信息
        res.errorCode = 1;
        res.errorMessage = "Can not open file";
        return res;
    }

    unsigned char head[12] = {0};
    file.read((char *)head, 8);
    std::string imgFormat = getFormatFromHeader(head);

    cv::Mat src_img = cv::imread(ImagePath);
    res = MattingRequestByImage(src_img, imgFormat);

    return res;
}

ImgBase64Result
MattingTritonClient::MattingRequestByBase64(const std::string &imageData) {
    ImgBase64Result res;

    if (imageData.empty()) {
        res.errorCode = 1;
        res.errorMessage = "base64_string size error";
        return res;
    }
    std::string decoded_data = base64_decode(imageData.c_str());
    unsigned char head[12] = {0};
    if (decoded_data.size() < 8) {
        res.errorCode = 1;
        res.errorMessage = "decoded_data size error";
        return res;
    }
    memcpy(head, decoded_data.data(), 8);
    std::string imgFormat = getFormatFromHeader(head);
    std::vector<uchar> data(decoded_data.begin(), decoded_data.end());
    cv::Mat Image = cv::imdecode(data, cv::IMREAD_COLOR);

    res = MattingRequestByImage(Image, imgFormat);

    return res;
}

std::string MattingTritonClient::GetConfigValue(const char *key,
                                                const char *engine) {
#ifdef TRITON_CLIENT
    std::string model_name = "model_bank";
    tc::Error err = client_->LoadModel(model_name);
    tc::InferOptions options(model_name);
    options.model_version_ = "";
    options.priority_ = 1;

    std::vector<std::string> input_key_data{std::string(key)};
    std::vector<std::string> input_engine_data{std::string(engine)};

    tc::InferInput *input0, *input1;
    std::vector<int64_t> shape{1};
    tc::InferInput::Create(&input0, "key", shape, "BYTES");
    tc::InferInput::Create(&input1, "engine", shape, "BYTES");
    std::shared_ptr<tc::InferInput> input0_ptr, input1_ptr;
    input0_ptr.reset(input0);
    input1_ptr.reset(input1);
    input0_ptr->AppendFromString(input_key_data);
    input1_ptr->AppendFromString(input_engine_data);
    std::vector<tc::InferInput *> inputs = {input0_ptr.get(), input1_ptr.get()};

    tc::InferRequestedOutput *output0;
    tc::Error err_out0 = tc::InferRequestedOutput::Create(&output0, "value");
    std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
    output0_ptr.reset(output0);
    std::vector<const tc::InferRequestedOutput *> outputs = {output0_ptr.get()};

    tc::InferResult *results;
    tc::Headers http_headers;
    tc::Error err_infer =
        client_->Infer(&results, options, inputs, outputs, http_headers);
    if (!err_infer.IsOk()) {
        std::cout << __LINE__ << " err_infer.Message(): " << err_infer.Message()
                  << std::endl;
        return "";
    }

    std::shared_ptr<tc::InferResult> results_ptr;
    results_ptr.reset(results);

    char *output0_data_ptr;
    size_t recv_output0_byte_size;

    std::vector<std::string> text_output_data;
    tc::Error err_paser = results_ptr->StringData("value", &text_output_data);
    if (!err_paser.IsOk()) {
        std::cout << __LINE__ << " err_paser.Message(): " << err_infer.Message()
                  << std::endl;
        return "";
    }
    return text_output_data[0];
#else
    return "";
#endif
}

ImgBase64Result
MattingTritonClient::MattingRequestByImage(cv::Mat Image,
                                           std::string imgFormat) {
    ImgBase64Result res = detect(Image, imgFormat);
    return res;
}

#ifdef TRITON_CLIENT
tc::Error MattingTritonClient::LoadModel() {
    model_name_ = GetConfigValue("matting_portrait", "matting_portrait");
    if (model_name_.size() < 1) {
        printf("line:%d model_name.size():%ld\n", __LINE__, model_name_.size());
        model_name_ = "matting_portrait";
    }
    tc::Error err = client_->LoadModel(model_name_.c_str());

    return err;
}

tc::Error MattingTritonClient::RemoveModel() {
    tc::Error err;

    return err;
}
#endif

ImgBase64Result MattingTritonClient::detect(cv::Mat &m_srcImg,
                                            std::string imgFormat) {
    ImgBase64Result res;
    res.errorCode = -1; // 假设成功时 errorCode 为 -1
    res.errorMessage.clear();

#ifdef TRITON_CLIENT
    cv::Mat transparentMat;
    cv::Mat result;
    cv::Mat image;
    uint b = 255, g = 255, r = 255, a = 0;

    if (m_srcImg.empty()) {
        res.errorCode = 2;
        res.errorMessage = "Image is empty";
        // 可以执行后续的图像处理操作
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage
                  << std::endl;
        return res;
    }
    // 检查图像是否为四通道,四通道转为三通道
    if (m_srcImg.channels() != 3) {
        cv::cvtColor(m_srcImg, m_srcImg, cv::COLOR_BGRA2BGR);
    }
    if (m_srcImg.channels() != 3) {
        res.errorCode = 2;
        res.errorMessage = "Please input RGB image.";
        // req error "unable to get unnorm_text_features"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage
                  << std::endl;
        return res;
    }

    // 图像预处理
    image = preprocess(m_srcImg);

    // 准备输入数据
    // Define the shape of the input tensor
    array<int64_t, 4> inputShape{1, 3, image.rows, image.cols};
    // Create a CPU memory allocator
    std::vector<int64_t> input_shape = {1, 3, inputHeight, inputWidth};

    std::vector<uint8_t> input_data(image.total() * image.channels() * 4);
    std::memcpy(input_data.data(), image.data, input_data.size());
    tc::InferInput *input0;

    tc::Error err_ino =
        tc::InferInput::Create(&input0, "input", input_shape, "FP32");
    if (!err_ino.IsOk()) {
        res.errorCode = 3;
        res.errorMessage = err_ino.Message();
        // model error "unable to get product_reviews"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage
                  << std::endl;
        return res;
    }

    std::shared_ptr<tc::InferInput> input0_ptr;
    input0_ptr.reset(input0);
    input0_ptr->SetBinaryData(true);
    input0_ptr->AppendRaw(input_data.data(), input_data.size());

    tc::InferRequestedOutput *output0;

    tc::Error err_out0 = tc::InferRequestedOutput::Create(&output0, "output");
    if (!err_out0.IsOk()) {
        res.errorCode = 4;
        res.errorMessage = err_out0.Message();
        // req error "unable to get unnorm_text_features"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage
                  << std::endl;
        return res;
    }

    std::shared_ptr<tc::InferRequestedOutput> output0_ptr;
    output0_ptr.reset(output0);

    std::string model_name = model_name_;
    std::string model_version = "";
    // The inference settings. Will be using default for now.
    tc::InferOptions options(model_name);
    options.model_version_ = model_version;
    options.priority_ = 1;

    std::vector<tc::InferInput *> inputs = {input0_ptr.get()};
    std::vector<const tc::InferRequestedOutput *> outputs = {output0_ptr.get()};

    tc::InferResult *results;
    tc::Headers http_headers;
    tc::Error err_infer =
        client_->Infer(&results, options, inputs, outputs, http_headers);
    if (!err_infer.IsOk()) {
        res.errorCode = 5;
        res.errorMessage = err_infer.Message();
        // Infer error "unable to run model"
        std::cout << __LINE__ << " res.errorMessage: " << res.errorMessage
                  << std::endl;
        return res;
    }
    std::shared_ptr<tc::InferResult> results_ptr;
    results_ptr.reset(results);

    float *output_data;
    size_t recv_output0_byte_size;
    results_ptr->RawData("output", (const uint8_t **)&output_data,
                         &recv_output0_byte_size);

    cv::Mat matte_np(inputHeight, inputWidth, CV_32F, output_data);

    // 调整灰度图大小与原始图像大小相匹配
    cv::Mat matte_resized;
    cv::resize(matte_np, matte_resized, cv::Size(m_srcImg.cols, m_srcImg.rows),
               0, 0, cv::INTER_LINEAR);
    matte_resized.convertTo(transparentMat, CV_8U,
                            255); // 将灰度图像转换为 8 位图像

    // 将抠图结果中背景部分设置为透明色
    result = cv::Mat(m_srcImg.size(), CV_8UC4);

    for (int y = 0; y < m_srcImg.rows; ++y) {
        for (int x = 0; x < m_srcImg.cols; ++x) {
            if (transparentMat.at<uchar>(y, x) >= 200) {
                cv::Vec3b bgr = m_srcImg.at<cv::Vec3b>(y, x);
                result.at<cv::Vec4b>(y, x) =
                    cv::Vec4b(bgr[0], bgr[1], bgr[2], 255); // 前景
            } else {
                result.at<cv::Vec4b>(y, x) = cv::Vec4b(b, g, r, a); // 背景
            }
        }
    }

    std::vector<uchar> buf;
    cv::imencode("." + imgFormat, result, buf);
    std::string res_base64 = base64_encode(buf.data(), buf.size());
    res.base64Result = res_base64;
    res.format = imgFormat; // 图片format
    res.errorCode = -1;     // 假设成功时 errorCode 为 -1
    res.errorMessage.clear();
#endif
    return res;
}
