/*
 * 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 <gtest/gtest.h>
#include <uuid.h>

#include <iostream>
#include <unordered_set>
#include <vector>

#include "constant.h"
#include "dataManagementVectorDatabase.h"
#include "embedding.h"

static std::string generateUUID() {
    uuid_t uuid;
    char uuidStr[37];
    uuid_generate_random(uuid);
    uuid_unparse(uuid, uuidStr);
    return uuidStr;
}

void replaceDashWithUnderscore(std::string& str) { std::replace(str.begin(), str.end(), '-', '_'); }

class DataManagementVectorDatabaseTest : public testing::Test {
public:
    static void SetUpTestSuite() { std::cout << "DataManagementVectorDatabaseTest SetUpTestSuite..." << std::endl; }

    static void TearDownTestSuite() {
        std::cout << "DataManagementVectorDatabaseTest TearDownTestSuite..." << std::endl;
    }

protected:
    void SetUp() override {
        std::cout << "DataManagementVectorDatabaseTest SetUp..." << std::endl;
        // 使用不同的集合名避免测试间干扰
        collectionName = "test_collection_" + generateUUID();
        replaceDashWithUnderscore(collectionName);  // 向量集合名之间要用下划线

        vectorDB = std::make_unique<DataManagementVectorDatabase>(collectionName);

        // 准备测试嵌入向量
        testEmbedding = Embedding(DataManagementEmbeddingDataType::Text, scene::TaskPriority::HIGH)
                            .getEmbedding("This is a test sentence for vector database.");
        if (testEmbedding.empty()) {
            FAIL() << "Failed to generate test embedding";
        }
    }

    void TearDown() override { std::cout << "DataManagementVectorDatabaseTest TearDown..." << std::endl; }

    std::unique_ptr<DataManagementVectorDatabase> vectorDB;
    std::string collectionName;
    std::vector<float> testEmbedding;
};

// TEST_F(DataManagementVectorDatabaseTest, test01) {
//     // 测试创建集合
//     EXPECT_TRUE(vectorDB != nullptr);

//     // 尝试插入数据验证集合是否存在
//     std::string uuid = generateUUID();
//     nlohmann::json metadata;
//     metadata["test_key"] = "test_value";

//     bool insertResult = vectorDB->insertVectorData({DataManagementVectorInfo{uuid, testEmbedding, metadata}});
//     EXPECT_TRUE(insertResult) << "Insert failed, collection might not exist";
//     // 清理测试数据
//     EXPECT_TRUE(vectorDB->deleteVectorData("id in ['" + uuid + "']"));
// }

// TEST_F(DataManagementVectorDatabaseTest, test02) {
//     // 测试插入与搜索功能
//     std::vector<DataManagementVectorInfo> testData;
//     const int testDataSize = 5;

//     for (int i = 0; i < testDataSize; ++i) {
//         std::string uuid = generateUUID();
//         nlohmann::json metadata;
//         metadata["index"] = i;
//         metadata["content"] = "Test content " + std::to_string(i);

//         // 对每条数据做轻微修改以增加多样性
//         std::vector<float> modifiedEmbedding = testEmbedding;
//         if (i > 0) {
//             modifiedEmbedding[0] += static_cast<float>(i) * 0.01f;
//         }

//         testData.emplace_back(DataManagementVectorInfo{uuid, modifiedEmbedding, metadata});
//     }

//     // 插入数据
//     EXPECT_TRUE(vectorDB->insertVectorData(testData));

//     // 搜索测试 - 使用第一条数据的嵌入向量作为查询
//     auto searchResult = vectorDB->searchMetadata(testData[0].embedding, testDataSize);

//     // 验证搜索结果
//     EXPECT_GE(searchResult.size(), 1) << "Should find at least the exact match";

//     // 验证第一条结果是否匹配（相似度应该很高,这里限制的是大于等于0.8）
//     EXPECT_GE(searchResult[0].similarity, 0.8f) << "Exact match similarity should be 0.8";
//     EXPECT_EQ(searchResult[0].metadata["index"], 0) << "Metadata should match inserted data";

//     // 验证结果顺序（相似度应该从高到低排序）
//     for (size_t i = 1; i < searchResult.size(); ++i) {
//         EXPECT_LE(searchResult[i].similarity, searchResult[i - 1].similarity)
//             << "Results should be sorted by similarity descending";
//     }

//     // 清理测试数据
//     std::string expression = "id in [";
//     for (const auto& data : testData) {
//         expression += "'" + data.id + "',";
//     }
//     expression.back() = ']';  // 替换最后一个逗号
//     EXPECT_TRUE(vectorDB->deleteVectorData(expression));
// }

// TEST_F(DataManagementVectorDatabaseTest, test03) {
//     // 测试搜索阈值
//     std::string uuid1 = generateUUID();
//     std::string uuid2 = generateUUID();

//     nlohmann::json metadata1;
//     metadata1["type"] = "type1";

//     nlohmann::json metadata2;
//     metadata2["type"] = "type2";

//     // 第二条数据的嵌入向量做较大修改
//     std::vector<float> embedding2 = testEmbedding;
//     for (auto& val : embedding2) {
//         val += 0.5f;  // 增加较大偏移
//     }

//     EXPECT_TRUE(vectorDB->insertVectorData({DataManagementVectorInfo{uuid1, testEmbedding, metadata1},
//                                             DataManagementVectorInfo{uuid2, embedding2, metadata2}}));

//     // 使用高阈值搜索 - 应该只返回非常相似的结果
//     auto highThresholdResult = vectorDB->searchMetadata(testEmbedding, 10, 0.8f);
//     EXPECT_EQ(highThresholdResult.size(), 1) << "High threshold should filter out dissimilar results";
//     EXPECT_EQ(highThresholdResult[0].id, uuid1) << "Should return only the exact match";

//     // 使用低阈值搜索 - 应该返回所有结果
//     auto lowThresholdResult = vectorDB->searchMetadata(testEmbedding, 10, 0.0f);
//     EXPECT_GE(lowThresholdResult.size(), 2) << "Low threshold should return all results";

//     // 清理
//     EXPECT_TRUE(vectorDB->deleteVectorData("id in ['" + uuid1 + "', '" + uuid2 + "']"));
// }

// TEST_F(DataManagementVectorDatabaseTest, test04) {
//     // 测试删除操作
//     std::string uuid1 = generateUUID();
//     std::string uuid2 = generateUUID();

//     nlohmann::json metadata;
//     metadata["test"] = "delete_test";

//     EXPECT_TRUE(vectorDB->insertVectorData({DataManagementVectorInfo{uuid1, testEmbedding, metadata},
//                                             DataManagementVectorInfo{uuid2, testEmbedding, metadata}}));

//     // 验证数据已插入
//     auto beforeDelete = vectorDB->searchMetadata(testEmbedding);
//     EXPECT_GE(beforeDelete.size(), 2);

//     // 删除第一条数据
//     EXPECT_TRUE(vectorDB->deleteVectorData("id == '" + uuid1 + "'"));

//     // 验证删除结果
//     auto afterDelete = vectorDB->searchMetadata(testEmbedding);
//     bool foundRemaining = false;
//     bool foundDeleted = false;

//     for (const auto& result : afterDelete) {
//         if (result.id == uuid1) foundDeleted = true;
//         if (result.id == uuid2) foundRemaining = true;
//     }

//     EXPECT_FALSE(foundDeleted) << "Deleted item should not appear in search results";
//     EXPECT_TRUE(foundRemaining) << "Other items should still exist";

//     // 清理剩余数据
//     EXPECT_TRUE(vectorDB->deleteVectorData("id == '" + uuid2 + "'"));
// }

// TEST_F(DataManagementVectorDatabaseTest, test05) {
//     // 测试插入 搜索空数据
//     auto emptyResult = vectorDB->searchMetadata({});
//     EXPECT_TRUE(emptyResult.empty()) << "Search with empty vector should return empty result";

//     // 测试删除不存在的数据
//     EXPECT_TRUE(vectorDB->deleteVectorData("id == 'non_existent_id'")) << "Delete non-existent data should not fail";

//     // 测试插入空向量
//     EXPECT_FALSE(vectorDB->insertVectorData({DataManagementVectorInfo{"", std::vector<float>{}, nlohmann::json{}}}))
//         << "Insert empty data should fail";

//     // 测试插入不匹配维度的向量
//     std::vector<float> wrongDimVector(10, 0.0f);
//     EXPECT_FALSE(
//         vectorDB->insertVectorData({DataManagementVectorInfo{generateUUID(), wrongDimVector, nlohmann::json{}}}))
//         << "Insert vector with wrong dimension should fail";
// }