#include "slotsnormalization.h"

#include <jsoncpp/json/json.h>

#include <filesystem>
#include <fstream>
#include <iostream>
#include <string>
#include <typeinfo>

#include "utils/logger.h"
#include "utils/util.h"
#include "variabletypeconversion.h"

using namespace VectorDB;

SlotsNormalization::SlotsNormalization() {
    client = Database::Create();
    /**
     * connect 有两种调用方式：
     * 1、传入ConnectParam配置安全信息等，参考 exampleSchema.cpp
     * 2、简化方式：无需传参，未设置安全信息
     */
    auto status = client->Connect();
    CheckStatus("Failed to connect milvus server:", status);
    LOGI("Connect to milvus server success.");
}

bool SlotsNormalization::CheckStatus(std::string &&prefix, const Status &status) {
    if (!status.IsOk()) {
        LOGI("!status.IsOk(): {} {} {} ", prefix, static_cast<int>(status.Code()), status.Message());
        return false;
    }
    return true;
}

bool SlotsNormalization::HasCollection() {
    bool exist;
    auto status = client->HasCollection(collectionName, exist);
    // todo 临时处理，向量数据库要修改 HasCollection 的返回
    // 如果返回报错也认为不存在，Failed to has TEST collection: 1002: collection
    // not found
    if (status.Code() == StatusCode::SERVER_FAILED) {
        return false;
    }
    CheckStatus("Failed to has TEST collection:", status);
    LOGI("Have TEST collection or not: {}", exist);

    return exist;
}

Json::Value SlotsNormalization::Search(const std::vector<float> &queryVector, const std::string &queryText,
                                       const std::string &slotName, const std::string &intentName) {
    // 根据向量化结果进行检索，语义搜索
    // 设置返回条数
    SearchArguments arguments(collectionName, filterNumber);
    // 查询输出哪些列
    arguments.AddOutputField(DEFAULT_ID_FIELD_NAME);
    arguments.AddOutputField(DYNAMIC_FIELD_NAME);
    arguments.AddOutputField(DEFAULT_VECTOR_FIELD_NAME);
    // 设置过滤条件
    std::string expr = "slot_name == '" + slotName + "' and intent_name == '" + intentName + "'";
    LOGI("expr: {}", expr);

    arguments.SetExpression(expr);

    arguments.SetGuaranteeTimestamp(GuaranteeStrongTs());

    arguments.AddTargetVector(DEFAULT_VECTOR_FIELD_NAME, std::move(queryVector));

    SearchResults search_results{};
    auto status = client->Search(arguments, search_results);
    bool requestSuccess = CheckStatus("Failed to search:", status);
    LOGI("\tSearch success. result size: {}", search_results.Results().size());

    Json::Value tempVectorSearch;

    if (requestSuccess) {
        tempVectorSearch["ToBeSearchedText"] = queryText;

        std::vector<Json::Value> similarityResult;
        std::vector<std::string> slotValue;

        for (auto &result : search_results.Results()) {
            auto &ids = result.Ids().IntIDArray();
            auto &distances = result.Scores();
            if (ids.size() != distances.size()) {
                LOGI("\tIllegal result!");
                continue;
            }

            if (ids.empty()) {
                LOGI("\tEmpty result!");
                continue;
            }

            auto idField = result.OutputField(DEFAULT_ID_FIELD_NAME);
            std::shared_ptr<Int64FieldData> idFieldPtr = std::static_pointer_cast<Int64FieldData>(idField);
            auto &idData = idFieldPtr->Data();

            auto jsonField = result.OutputField(DYNAMIC_FIELD_NAME);
            std::shared_ptr<JsonFieldData> jsonFieldPtr = std::static_pointer_cast<JsonFieldData>(jsonField);
            auto &jsonData = jsonFieldPtr->Data();

            auto vectorField = result.OutputField(DEFAULT_VECTOR_FIELD_NAME);
            std::shared_ptr<FloatVecFieldData> vectorFieldPtr =
                std::static_pointer_cast<FloatVecFieldData>(vectorField);
            auto &vectorData = vectorFieldPtr->Data();

            for (size_t i = 0; i < ids.size(); ++i) {
                LOGI("\tID: {} \tSimilarity: {} \t{}: {} \t{}: {}", ids[i], distances[i], idFieldPtr->Name(), idData[i],
                     jsonFieldPtr->Name(), jsonData[i].dump());

                Json::Value tempJsonData = util::stringToJson(jsonData[i].dump());

                if (distances[i] >= similarityThreshold) {
                    Json::Value distanceRes;
                    distanceRes["slotValue"] = tempJsonData["slot_value"].asString();
                    distanceRes["slotName"] = tempJsonData["slot_name"].asString();
                    distanceRes["intentName"] = tempJsonData["intent_name"].asString();
                    distanceRes["Distance"] = distances[i];

                    similarityResult.push_back(distanceRes);
                    slotValue.push_back(tempJsonData["slot_value"].asString());
                }
            }
        }

        Json::Value similarityArray(Json::arrayValue);
        for (const auto &value : similarityResult) {
            similarityArray.append(value);
        }
        tempVectorSearch["Top_K_Results"] = similarityArray;

        Json::Value slotValueArray(Json::arrayValue);
        for (const auto &value : slotValue) {
            slotValueArray.append(value);
        }
        tempVectorSearch["Top_K_Label"] = slotValueArray;
    }

    return tempVectorSearch;
}

SlotsNormalization::~SlotsNormalization() {
    auto status = client->Disconnect();
    CheckStatus("Failed to Disconnect milvus server:", status);
    LOGI("Successfully disconnected from milvus server.");
}
