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

#include <unistd.h>
#include <getopt.h>

#include <cassert>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <nlohmann/json.hpp>

#include "utils.h"

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

namespace ai_engine::core_ai::embedding {
EmbeddingProxy::EmbeddingProxy() {}

bool EmbeddingProxy::createTritonClient() {
#ifdef TRITON_CLIENT
    if (initTritonClient_)
        return true;
    tc::Error error = tc::InferenceServerGrpcClient::Create(
        &grpcClient_, grpcServerUrl_, true /* verbose */);

    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Create http client error,"
                  << error.Message() << std::endl;
        return false;
    }
    initTritonClient_ = true;
#endif
    return true;
}
EmbeddingProxyResult EmbeddingProxy::loadTextModel() {
#ifdef TRITON_CLIENT
    if (initTextModel_)
        return EmbeddingProxyResult(true, std::vector<float>{});
    tc::Error error = grpcClient_->LoadModel(textEmbeddingModel_);

    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Load text-embedding-text-model: "
                  << textEmbeddingModel_ << " error"
                  << " error message: " << error.Message() << std::endl;
        return EmbeddingProxyResult(false, LOAD_MODEL_ERROR, error.Message());
    }
    initTextModel_ = true;
#endif
    return EmbeddingProxyResult(true, std::vector<float>{});
}
EmbeddingProxyResult EmbeddingProxy::loadImageModel() {
#ifdef TRITON_CLIENT
    if (initImageModel_)
        return EmbeddingProxyResult(true, std::vector<float>{});

    tc::Error error = grpcClient_->LoadModel(textEmbeddingImageModel_);

    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Load text-embedding-image-model:"
                  << textEmbeddingImageModel_ << " error"
                  << " error message: " << error.Message() << std::endl;
        return EmbeddingProxyResult(false, LOAD_MODEL_ERROR, error.Message());
    }
    error = grpcClient_->LoadModel(imageEmbeddingModel_);
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Load image-embedding-image-model: "
                  << imageEmbeddingModel_ << " error"
                  << " error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, LOAD_MODEL_ERROR, error.Message());
    }
    initImageModel_ = true;
#endif
    return EmbeddingProxyResult(true, std::vector<float>{});
}
std::string EmbeddingProxy::modelInfo() const {
    const std::string modelInfo = R"(
            {"models":{
                "text_model":{
                            "name":"ensemble-embd_gte-base_uint8-text", 
                            "dim":768,
                            "ondevice":true},
                "image_model":{
                            "image_text_name":"ensemble-embd_cn-clip_512-uint8-text",
                            "image_name":"embd_cn-clip_512-uint8-image",
                            "dim":1024,
                            "ondevice":true}
                }
            })";
    return modelInfo;
}
EmbeddingProxyResult EmbeddingProxy::embeddingText(const std::string &text) {
#ifdef TRITON_CLIENT
    std::vector<std::string> inputData(1);
    inputData[0] = text;

    std::vector<int64_t> shape{1};

    // Initialize the inputs with the data.
    tc::InferInput *input0;
    tc::Error error =
        tc::InferInput::Create(&input0, "product_reviews", shape, "BYTES");
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error: unable to get product_reviews,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferInput> inputPtr;
    inputPtr.reset(input0);
    inputPtr->SetBinaryData(true);

    error = inputPtr->AppendFromString(inputData);
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to set data for "
                     "product_reviews,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    // Generate the outputs to be requested.
    tc::InferRequestedOutput *output0;

    error = tc::InferRequestedOutput::Create(&output0, "logits");
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to get logits,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferRequestedOutput> outputPtr;
    outputPtr.reset(output0);
    outputPtr->SetBinaryData(true);

    // The inference settings. Will be using default for now.
    tc::InferOptions options(textEmbeddingModel_);
    options.model_version_ = textEmbeddingVersion_;

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

    tc::InferResult *results;
    error = grpcClient_->Infer(&results, options, inputs, outputs);
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton infer error: unable to run model,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferResult> resultsPtr;
    resultsPtr.reset(results);

    // Validate the results...
    if (!validateShapeAndDatatype("logits", resultsPtr, "FP32", 0)) {
        std::cerr << "[Embedding proxy]: Trion infer error: embeddingText "
                     "validateShapeAndDatatype failed."
                  << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR,
                                    "Trion infer error: embeddingText "
                                    "validateShapeAndDatatype failed.");
    }

    std::vector<int64_t> shapeOut;
    resultsPtr->Shape("logits", &shapeOut);

    float *outputDataPtr;

    const size_t outputByteSize =
        textModelDim_ * 4 * shapeOut[1] *
        shapeOut[0]; // 输出 1*756*768  会输出756(会变化)个768，取第一个768
                     // 所以是 768 * 4 * shapeOut[1] 字节

    size_t recvOutputByteSize;
    error = resultsPtr->RawData("logits", (const uint8_t **)&outputDataPtr,
                                &recvOutputByteSize);
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error: unable to get data for logits,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    if (recvOutputByteSize != outputByteSize) {
        std::cerr << "[Embedding proxy]:Error: received incorrect byte size "
                     "for 'logits': "
                  << recvOutputByteSize << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR,
                                    "Trion infer error: embeddingText received "
                                    "incorrect byte size.");
    }
    std::vector<float> tokens;

    for (size_t i = 0; i < textModelDim_; ++i) {
        tokens.push_back(outputDataPtr[i]);
    }

    normalize(tokens);

    return EmbeddingProxyResult(true, tokens);
#else
    return EmbeddingProxyResult(true, {});
#endif
}
EmbeddingProxyResult EmbeddingProxy::embeddingImage(const std::string &fileData,
                                                    int fileDataType) {
#ifdef TRITON_CLIENT
    cv::Mat image;
    try {
        if (fileDataType == 0)
            image = cv::imread(fileData, cv::IMREAD_COLOR);
        else if (fileDataType == 1) {
            std::vector<unsigned char> decodedImageData =
                base64Decode(fileData);
            image = cv::imdecode(decodedImageData, cv::IMREAD_COLOR);
        }
    } catch (const cv::Exception &e) {
        std::cerr << "[Embedding proxy]:OpenCV Error:" << e.what() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR,
                                    "Trion infer error: embeddingImage "
                                    "preprocessed image failed.");
    }

    if (image.empty()) {
        if (fileDataType == 0) {
            // 输入文件路径的才打印，否则base64数据太多了
            std::cerr << "[Embedding proxy]:Error:Unable to load image from:"
                      << fileData << std::endl;
        }
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR,
                                    "Trion infer error: embeddingImage "
                                    "preprocessed image is empty.");
    }
    std::vector<float> input0_data = preprocessVision(image);

    std::vector<int64_t> shape = {3, 224, 224};

    // Initialize the inputs with the data.
    tc::InferInput *input0;
    tc::Error error = tc::InferInput::Create(&input0, "image", shape, "FP32");
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error:unable to get image,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    std::shared_ptr<tc::InferInput> inputPtr;
    inputPtr.reset(input0);
    error = inputPtr->AppendRaw(reinterpret_cast<uint8_t *>(&input0_data[0]),
                                input0_data.size() * sizeof(float));
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error:unable to set data for image,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    // The inference settings. Will be using default for now.
    tc::InferOptions options(imageEmbeddingModel_);
    options.model_version_ = imageEmbeddingVersion_;

    std::vector<tc::InferInput *> inputs = {inputPtr.get()};
    // Empty output vector will request data for all the output tensors from
    // the server.
    std::vector<const tc::InferRequestedOutput *> outputs = {};

    tc::InferResult *results;
    error = grpcClient_->Infer(&results, options, inputs, outputs);
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error:unable to run model,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferResult> resultsPtr;
    resultsPtr.reset(results);

    // Validate the results...
    if (!validateShapeAndDatatype("unnorm_image_features", resultsPtr, "FP32",
                                  1)) {
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR,
                                    "Trion infer error: embeddingImage "
                                    "validateShapeAndDatatype failed.");
    }

    // Get pointers to the result returned...
    float *outputData;
    size_t outputByteSize;
    error = resultsPtr->RawData("unnorm_image_features",
                                (const uint8_t **)&outputData, &outputByteSize);
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error:unable to get result data for "
               "unnorm_image_features,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    if (outputByteSize != imageModelDim_ * sizeof(float)) {
        std::cerr << "[Embedding proxy]:Triton error:received incorrect byte "
                     "size for "
                     "'unnorm_image_features':"
                  << outputByteSize << std::endl;
        return EmbeddingProxyResult(
            false, TRITON_INFER_ERROR,
            "Trion infer error: embeddingImage received "
            "incorrect byte size.");
    }

    std::vector<float> output0_data_vector(outputData,
                                           outputData + imageModelDim_);
    normalize(output0_data_vector);

    return EmbeddingProxyResult{true, output0_data_vector};
#else
    return EmbeddingProxyResult{true, {}};
#endif
}

EmbeddingProxyResult
EmbeddingProxy::embeddingTextWithImageModel(const std::string &text) {
#ifdef TRITON_CLIENT
    std::vector<std::string> input0_data(1);
    input0_data[0] = text;

    std::vector<int64_t> shape{1};

    // Initialize the inputs with the data.
    tc::InferInput *input0;
    tc::Error error =
        tc::InferInput::Create(&input0, "product_reviews", shape, "BYTES");
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error: unable to get product_reviews,"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferInput> inputPtr;
    inputPtr.reset(input0);
    inputPtr->SetBinaryData(true);
    error = inputPtr->AppendFromString(input0_data);
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to set data for "
                     "product_reviews,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    // Generate the outputs to be requested.
    tc::InferRequestedOutput *output0;
    error = tc::InferRequestedOutput::Create(&output0, "unnorm_text_features");
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to get "
                     "unnorm_text_features,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }
    std::shared_ptr<tc::InferRequestedOutput> outputPtr;
    outputPtr.reset(output0);
    outputPtr->SetBinaryData(true);

    // The inference settings. Will be using default for now.
    tc::InferOptions options(textEmbeddingImageModel_);
    options.model_version_ = textEmbeddingImageVersion_;
    // options.priority_ = 3;

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

    tc::InferResult *results;
    error = grpcClient_->Infer(&results, options, inputs, outputs);
    if (!error.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to run model,"
                  << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

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

    // Validate the results...
    if (!validateShapeAndDatatype("unnorm_text_features", resultsPtr, "FP32",
                                  1)) {
        return EmbeddingProxyResult(
            false, TRITON_INFER_ERROR,
            "Trion infer error: embeddingTextWithImageModel "
            "validateShapeAndDatatype failed.");
    }

    float *outputDataPtr;

    size_t outputByteSize =
        imageModelDim_ * 4; // 输出 1024 个 float 所以是 1024 * 4 字节

    size_t recvOutputByteSize;
    error = resultsPtr->RawData("unnorm_text_features",
                                (const uint8_t **)&outputDataPtr,
                                &recvOutputByteSize);
    if (!error.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error: unable to get result data for "
               "unnorm_text_features"
            << "error message:" << error.Message() << std::endl;
        return EmbeddingProxyResult(false, TRITON_INFER_ERROR, error.Message());
    }

    if (recvOutputByteSize != outputByteSize) {
        std::cerr
            << "[Embedding proxy]:Error: received incorrect byte size for "
               "'unnorm_text_features': "
            << recvOutputByteSize << std::endl;
        return EmbeddingProxyResult(
            false, TRITON_INFER_ERROR,
            "Trion infer error: embeddingTextWithImageModel received "
            "incorrect byte size");
    }
    std::vector<float> tokens;

    //   std::cout << "unnorm_text_features: ";
    for (size_t i = 0; i < imageModelDim_; ++i) {
        // std::cout << outputDataPtr[i] << " ";
        tokens.push_back(outputDataPtr[i]);
    }
    //   std::cout << std::endl;

    normalize(tokens);
    return EmbeddingProxyResult{true, tokens};
#else
    return EmbeddingProxyResult{true, {}};

#endif
}

#ifdef TRITON_CLIENT
bool EmbeddingProxy::validateShapeAndDatatype(
    const std::string &name, std::shared_ptr<tc::InferResult> result,
    const std::string &expectedDataType, const int modelSwitch) {
    std::vector<int64_t> shape;
    tc::Error err;
    err = result->Shape(name, &shape);
    if (!err.IsOk()) {
        std::cerr << "[Embedding proxy]:Triton error: unable to get shape for"
                  << name << std::endl;
        return false;
    }
    // Validate shape
    if (modelSwitch == 0) {
        if ((shape.size() != 3) || (shape[2] != textModelDim_)) {
            std::cerr
                << "[Embedding proxy]:Error: received incorrect shapes for "
                << name << std::endl;
            return false;
        }
    } else if (modelSwitch == 1) {
        if (shape != std::vector<int64_t>{imageModelDim_}) {
            std::cerr
                << "[Embedding proxy]:Error: received incorrect shapes for "
                << name << std::endl;
            return false;
        }
    }

    std::string dataType;
    err = result->Datatype(name, &dataType);
    if (!err.IsOk()) {
        std::cerr
            << "[Embedding proxy]:Triton error: unable to get dataType for"
            << name << std::endl;
        return false;
    }

    // Validate dataType
    if (dataType.compare(expectedDataType) != 0) {
        std::cerr << "[Embedding proxy]:Error: received incorrect dataType for "
                  << name << ": " << dataType << std::endl;
        return false;
    }
    return true;
}
#endif
std::vector<unsigned char>
EmbeddingProxy::base64Decode(const std::string &input) {
    std::vector<unsigned char> output;
    const std::string base64Chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                                    "abcdefghijklmnopqrstuvwxyz"
                                    "0123456789+/";

    int val = 0, valb = -8;
    for (auto c : input) {
        if (std::isspace(c)) {
            continue;
        }
        if (c == '=') {
            break;
        }
        const auto ch = std::find(base64Chars.begin(), base64Chars.end(), c);
        if (ch == base64Chars.end()) {
            std::cerr << "[Embedding proxy]:parse base64Decode error!"
                      << std::endl;
            throw std::runtime_error("Invalid base64 character");
        }

        val = (val << 6) + (ch - base64Chars.begin());
        valb += 6;
        if (valb >= 0) {
            output.push_back(static_cast<unsigned char>((val >> valb) & 0xFF));
            valb -= 8;
        }
    }

    return output;
}
std::vector<float> EmbeddingProxy::imageToFloatVector(const cv::Mat &image) {
    std::vector<float> floatVector;

    // 确保图像不为空且为三通道图像
    if (image.empty() || image.channels() != 3) {
        std::cerr << "[Embedding proxy]:imageToFloatVector error:Image is "
                     "empty or not 3 channels, "
                     "can't convert to vector"
                  << std::endl;
        return floatVector;
    }

    // 将每个通道的像素值添加到向量中
    for (int c = 0; c < image.channels(); ++c) {
        for (int i = 0; i < image.cols; ++i) {
            for (int j = 0; j < image.rows; ++j) {
                float pixelValue = image.at<cv::Vec3f>(i, j)[c];
                floatVector.push_back(pixelValue);
            }
        }
    }

    return floatVector;
}
std::vector<float> EmbeddingProxy::preprocessVision(cv::Mat &image) {
    const int width = 224;
    const int height = 224;
    const cv::Scalar mean = cv::Scalar(0.48145466, 0.4578275, 0.40821073);
    const cv::Scalar deviation = cv::Scalar(0.26862954, 0.26130258, 0.27577711);

    cv::resize(image, image, cv::Size(width, height), cv::INTER_CUBIC);
    cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

    // Normalize
    image.convertTo(image, CV_32FC3, 1.0 / 255.0);
    cv::subtract(image, mean, image);
    cv::divide(image, deviation, image);

    return imageToFloatVector(image);
}
} // namespace ai_engine::core_ai::embedding