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

#include <log.h>
#include <rpc/vectorDB/vectorDatabase.h>

DataManagementVectorDatabase::DataManagementVectorDatabase(const std::string& collectionName, int dimension)
    : collectionName_(collectionName), dimension_(dimension) {
    if (!scene::VectorDatabase::getInstance().hasCollection(collectionName_)) {
        PRINT_DEBUG("DataManagementVectorDatabase create collection: %s.\n", collectionName_.c_str());
        bool ret = createCollection();
        if (!ret) {
            PRINT_ERROR("createCollection %s failed.\n", collectionName_.c_str());
            return;
        }
    }
}

bool DataManagementVectorDatabase::insertVectorData(const std::vector<DataManagementVectorInfo>& vectorInfos) {
    std::vector<std::string> idVec;
    idVec.reserve(vectorInfos.size());
    std::vector<std::vector<float>> floatVec;
    floatVec.reserve(vectorInfos.size());
    std::vector<nlohmann::json> metaVec;
    metaVec.reserve(vectorInfos.size());

    for (const auto& info : vectorInfos) {
        idVec.emplace_back(info.id);
        floatVec.emplace_back(info.embedding);
        metaVec.emplace_back(info.metadata);
    }

    std::vector<VectorDB::FieldDataPtr> data{
        std::make_shared<VectorDB::VarCharFieldData>(DEFAULT_ID_FIELD_NAME, idVec),
        std::make_shared<VectorDB::FloatVecFieldData>(DEFAULT_VECTOR_FIELD_NAME, floatVec),
        std::make_shared<VectorDB::JsonFieldData>(DYNAMIC_FIELD_NAME, metaVec)};
    VectorDB::DmlResults results;
    int ret = scene::VectorDatabase::getInstance().insert(collectionName_, data, results);
    PRINT_DEBUG("%s vector insert rows: %zu\n", collectionName_.c_str(), results.IdArray().StrIDArray().size());
    return ret;
}

bool DataManagementVectorDatabase::deleteVectorData(const std::string& expression) {
    VectorDB::DmlResults results;
    bool ret = scene::VectorDatabase::getInstance().deleteData(collectionName_, expression, results);
    if (!ret) {
        PRINT_ERROR("%s delete vector db failed. expression: %s.\n", collectionName_.c_str(), expression.c_str());
        return false;
    }
    PRINT_DEBUG("%s vector delete result: %zu\n", collectionName_.c_str(), results.IdArray().StrIDArray().size());
    return true;
}

SimilaritySearchMetadataResult DataManagementVectorDatabase::searchMetadata(const std::vector<float>& embedding,
                                                                            int topK, float threshold) {
    VectorDB::SearchArguments arguments(collectionName_, topK);
    arguments.AddOutputField(DYNAMIC_FIELD_NAME);  // 查询只输出 metadata 列
    arguments.SetGuaranteeTimestamp(VectorDB::GuaranteeStrongTs());
    arguments.AddTargetVector(DEFAULT_VECTOR_FIELD_NAME, embedding);

    VectorDB::SearchResults results;
    bool ret = scene::VectorDatabase::getInstance().search(arguments, results);
    if (!ret) {
        PRINT_ERROR("SearchMetadata %s failed. embedding size: %zu\n", collectionName_.c_str(), embedding.size());
        return SimilaritySearchMetadataResult();
    }

    SimilaritySearchMetadataResult output;
    for (const auto& result : results.Results()) {
        auto& ids = result.Ids().StrIDArray();
        if (ids.empty()) {
            PRINT_INFO("SearchMetadata %s result is empty.\n", collectionName_.c_str());
            continue;
        }

        auto similarity = result.Scores();
        if (ids.size() != similarity.size()) {
            PRINT_ERROR("Illegal %s result! %zu vs %zu\n", collectionName_.c_str(), ids.size(), similarity.size());
            continue;
        }

        auto dynamicField = result.OutputField(DYNAMIC_FIELD_NAME);
        VectorDB::JsonFieldDataPtr dynamicFieldPtr = std::static_pointer_cast<VectorDB::JsonFieldData>(dynamicField);
        auto& dynamicFieldData = dynamicFieldPtr->Data();

        for (size_t i = 0; i < ids.size(); ++i) {
            // 根据相似度阈值筛选结果，数据库查询结果是按照相似度倒排的，如果出现相似度小于阈值的，后面的结果都比阈值小
            if (similarity[i] < threshold) {
                PRINT_INFO("SearchMetadata similarity %f less than threshold %f.\n", similarity[i], threshold);
                break;
            }

            PRINT_DEBUG("SearchMetadata id: %s, dynamic data: %s, similarity: %f.\n", ids[i].c_str(),
                        dynamicFieldData[i].dump().c_str(), similarity[i]);
            output.emplace_back(MetadataResult{ids[i], similarity[i], dynamicFieldData[i]});
        }
    }
    return output;
}

bool DataManagementVectorDatabase::createCollection() {
    VectorDB::CollectionSchema schema(collectionName_, collectionName_, true);
    schema.AddField(
        VectorDB::FieldSchema(DEFAULT_ID_FIELD_NAME, VectorDB::DataType::VARCHAR, "id", true).WithMaxLength(50));
    schema.AddField(VectorDB::FieldSchema(DEFAULT_VECTOR_FIELD_NAME, VectorDB::DataType::FLOAT_VECTOR, "vector")
                        .WithDimension(dimension_));

    // create index
    VectorDB::IndexDesc desc(DEFAULT_VECTOR_FIELD_NAME, collectionName_ + "_vector_index", 0);

    return scene::VectorDatabase::getInstance().createCollection(schema, desc);
}
