#include "gtest/gtest.h"
#include "db/NewDB.h"  // NewDB 的头文件
#include "leveldb/env.h"
#include "leveldb/db.h"
#include "db/write_batch_internal.h"
#include <iostream>
#include <random>
#include <ctime>
#include <thread>

using namespace leveldb;

Status OpenNewDB(std::string dbName, NewDB** db) {
    Options options = Options();
    options.create_if_missing = true;
    return NewDB::Open(options, dbName, db);
}

Status OpenDB(std::string dbName, DB **db) {
  Options options;
  options.create_if_missing = true;
  return DB::Open(options, dbName, db);
}

// 全局的随机数引擎
std::default_random_engine rng;

// 设置随机种子
void SetGlobalSeed(unsigned seed) {
    rng.seed(seed);
}

// 生成随机字符串
std::string GenerateRandomString(size_t length) {
    static const char alphanum[] =
        "0123456789"
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz";

    std::uniform_int_distribution<int> dist(0, sizeof(alphanum) - 2);

    std::string str(length, 0);
    for (size_t i = 0; i < length; ++i) {
        str[i] = alphanum[dist(rng)];
    }
    return str;
}

std::string testdbname = "dbtest76";


TEST(TestNewDB, PUTRecoveryTest) {
// put_fields-新增
    // 创建 NewDB 实例
    NewDB* db;
    ASSERT_TRUE(OpenNewDB(testdbname, &db).ok());

    db->CreateIndexOnField("address");
    // 批量插入数据
    std::vector<std::pair<std::string, FieldArray>> bulk_data;
    std::vector<std::string> bulk_index;
    int num_records = 100;
    for (int i = 0; i < num_records; ++i) {
        std::string key = "k_" + GenerateRandomString(10);
        FieldArray fields = {
            {"name", GenerateRandomString(5)},
            {"address", GenerateRandomString(15)},
            {"phone", GenerateRandomString(11)}
        };
        bulk_data.emplace_back(key, fields);
        bulk_index.emplace_back(db->ConstructIndexKey(key, fields[1]));
        
        // std::cout << key << "..." << db->SerializeValue(fields) << std::endl;

        ASSERT_TRUE(db->Put_fields(WriteOptions(), key, fields).ok());
    }
    // for(auto& key:bulk_index){
    //     std::cout << key << std::endl;
    // }
    delete db;

// 模拟崩溃
    // 创建 DB 实例
    DB* data_db;
    ASSERT_TRUE(OpenDB(testdbname + "_data", &data_db).ok());
    // 删除部分数据
    int num_updates = 20;
    for (int i = 0; i < num_updates; ++i) {
        auto& [key, fields] = bulk_data[i];
        ASSERT_TRUE(data_db->Delete(WriteOptions(), key).ok());
    }
    delete data_db;

    // 创建 DB 实例
    DB* index_db;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db).ok());
    for (int i = 0; i < num_updates; ++i) {
        auto& key = bulk_index[i];      
        ASSERT_TRUE(index_db->Delete(WriteOptions(), key).ok());
    }
    delete index_db;


// 测试是否能通过recoverDB恢复
    // 创建 NewDB 实例
    NewDB* db_2;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_2).ok());

    // 验证插入的数据和索引
    for (const auto& [key, fields] : bulk_data) {
        // 获取并验证字段数据
        FieldArray retrieved_fields;
        ASSERT_TRUE(db_2->Get_fields(ReadOptions(), Slice(key), &retrieved_fields).ok());

        // 验证插入的每个字段
        for (size_t i = 0; i < fields.size(); ++i) {
            EXPECT_EQ(fields[i].first, retrieved_fields[i].first);  // 字段名
            EXPECT_EQ(fields[i].second, retrieved_fields[i].second); // 字段值
        }
    }
    delete db_2;

    // 创建 DB 实例
    DB* index_db_2;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db_2).ok());
    for (const auto& key : bulk_index) {
        // 获取并验证字段数据
        std::string value;
        ASSERT_TRUE(index_db_2->Get(ReadOptions(), Slice(key), &value).ok());
    }
    delete index_db_2;
}

TEST(TestNewDB, UpdateRecoveryTest) {
// put_fields-更新
    // 创建 NewDB 实例
    NewDB* db_3;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_3).ok());

    db_3->CreateIndexOnField("address");
    // 批量插入数据
    std::vector<std::pair<std::string, FieldArray>> bulk_data;
    std::vector<std::string> bulk_index;
    int num_records = 100;
    for (int i = 0; i < num_records; ++i) {
        std::string key = "k_" + GenerateRandomString(10);
        FieldArray fields = {
            {"name", GenerateRandomString(5)},
            {"address", GenerateRandomString(15)},
            {"phone", GenerateRandomString(11)}
        };
        bulk_data.emplace_back(key, fields);
        bulk_index.emplace_back(db_3->ConstructIndexKey(key, fields[1]));
        
        // std::cout << key << "..." << db->SerializeValue(fields) << std::endl;

        ASSERT_TRUE(db_3->Put_fields(WriteOptions(), key, fields).ok());
    }

    // 更新数据
    std::vector<std::pair<std::string, FieldArray>> bulk_data_update;
    std::vector<std::string> bulk_index_update;
    num_records = 20;
    for (int i = 0; i < num_records; ++i) {
        auto& [key, fields] = bulk_data[i];
        FieldArray fields_update = {
            {"name", fields[0].second},
            {"address", GenerateRandomString(15)},
            {"phone", fields[2].second}
        };
        bulk_data_update.emplace_back(key, fields_update);
        bulk_index_update.emplace_back(db_3->ConstructIndexKey(key, fields_update[1]));
        
        ASSERT_TRUE(db_3->Put_fields(WriteOptions(), key, fields_update).ok());
    }

    delete db_3;


// 模拟崩溃
    // 创建 DB 实例
    DB* data_db_2;
    ASSERT_TRUE(OpenDB(testdbname + "_data", &data_db_2).ok());
    // 删除部分数据
    int num_updates = 10;
    for (int i = 0; i < num_updates; ++i) {
        auto& [key, fields] = bulk_data_update[i];
        ASSERT_TRUE(data_db_2->Delete(WriteOptions(), key).ok());
    }
    delete data_db_2;

    // 创建 DB 实例
    DB* index_db_3;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db_3).ok());
    for (int i = 0; i < num_updates; ++i) {
        auto& key = bulk_index_update[i];      
        ASSERT_TRUE(index_db_3->Delete(WriteOptions(), key).ok());
    }
    delete index_db_3;

// 测试是否能通过recoverDB恢复
    // 创建 NewDB 实例
    NewDB* db_4;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_4).ok());

    // 验证插入的数据和索引
    for (const auto& [key, fields] : bulk_data_update) {
        // 获取并验证字段数据
        FieldArray retrieved_fields;
        ASSERT_TRUE(db_4->Get_fields(ReadOptions(), Slice(key), &retrieved_fields).ok());

        // 验证插入的每个字段
        for (size_t i = 0; i < fields.size(); ++i) {
            EXPECT_EQ(fields[i].first, retrieved_fields[i].first);  // 字段名
            EXPECT_EQ(fields[i].second, retrieved_fields[i].second); // 字段值
        }
    }
    delete db_4;

    // 创建 DB 实例
    DB* index_db_4;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db_4).ok());
    for (const auto& key : bulk_index_update) {
        // 获取并验证字段数据
        std::string value;
        ASSERT_TRUE(index_db_4->Get(ReadOptions(), Slice(key), &value).ok());
    }
    delete index_db_4;
}

TEST(TestNewDB, DeleteRecoveryTest) {
    // 创建 NewDB 实例
    NewDB* db_3;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_3).ok());

    db_3->CreateIndexOnField("address");
    // 批量插入数据
    std::vector<std::pair<std::string, FieldArray>> bulk_data;
    std::vector<std::string> bulk_index;
    int num_records = 100;
    for (int i = 0; i < num_records; ++i) {
        std::string key = "k_" + GenerateRandomString(10);
        FieldArray fields = {
            {"name", GenerateRandomString(5)},
            {"address", GenerateRandomString(15)},
            {"phone", GenerateRandomString(11)}
        };
        bulk_data.emplace_back(key, fields);
        bulk_index.emplace_back(db_3->ConstructIndexKey(key, fields[1]));
    
        ASSERT_TRUE(db_3->Put_fields(WriteOptions(), key, fields).ok());
    }

    // 删除数据
    std::vector<std::pair<std::string, FieldArray>> bulk_data_delete;
    std::vector<std::string> bulk_index_delete;
    num_records = 20;
    for (int i = 0; i < num_records; ++i) {
        auto& [key, fields] = bulk_data[i];
        bulk_data_delete.emplace_back(key, fields);
        bulk_index_delete.emplace_back(db_3->ConstructIndexKey(key, fields[1]));
        
        ASSERT_TRUE(db_3->Delete(WriteOptions(), key));
    }

    delete db_3;

// 模拟崩溃
    // 创建 DB 实例
    DB* data_db_2;
    ASSERT_TRUE(OpenDB(testdbname + "_data", &data_db_2).ok());
    // 恢复部分数据
    int num_updates = 10;
    for (int i = 0; i < num_updates; ++i) {
        auto& [key, fields] = bulk_data_delete[i];
        ASSERT_TRUE(data_db_2->Put_fields(WriteOptions(), key, fields).ok());
    }
    delete data_db_2;

    // 创建 DB 实例
    DB* index_db_3;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db_3).ok());
    for (int i = 0; i < num_updates; ++i) {
        auto& key = bulk_index_delete[i];      
        ASSERT_TRUE(index_db_3->Put(WriteOptions(), key, "").ok());
    }
    delete index_db_3;

// 测试是否能通过recoverDB恢复
    // 创建 NewDB 实例
    NewDB* db_4;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_4).ok());

    // 验证插入的数据和索引
    for (const auto& [key, fields] : bulk_data_delete) {
        // 获取并验证字段数据
        FieldArray retrieved_fields;
        ASSERT_FALSE(db_4->Get_fields(ReadOptions(), Slice(key), &retrieved_fields).ok());
    }
    delete db_4;

    // 创建 DB 实例
    DB* index_db_4;
    ASSERT_TRUE(OpenDB(testdbname + "_index", &index_db_4).ok());
    for (const auto& key : bulk_index_delete) {
        // 获取并验证字段数据
        std::string value;
        ASSERT_FALSE(index_db_4->Get(ReadOptions(), Slice(key), &value).ok());
    }
    delete index_db_4;
}


TEST(TestNewDB, QueryIndexTest) {
    // 创建 NewDB 实例
    NewDB* db;
    ASSERT_TRUE(OpenNewDB(testdbname, &db).ok());

    db->CreateIndexOnField("address");

    std::string key = "k_" + GenerateRandomString(10);
    FieldArray fields = {
        {"name", GenerateRandomString(5)},
        {"address", GenerateRandomString(15)},
        {"phone", GenerateRandomString(11)}
    };
    ASSERT_TRUE(db->Put_fields(WriteOptions(), key, fields).ok());

    delete db;

    // 创建 DB 实例
    DB* data_db;
    ASSERT_TRUE(OpenDB(testdbname + "_data", &data_db).ok());
    // update the address
    FieldArray fields_new = {
        {"name", fields[0].second},
        {"address", GenerateRandomString(15)},
        {"phone", fields[2].second}
    };
    ASSERT_TRUE(data_db->Put_fields(WriteOptions(), key, fields_new).ok());
    delete data_db; 

    // 创建 NewDB 实例
    NewDB* db_2;
    ASSERT_TRUE(OpenNewDB(testdbname, &db_2).ok());

    // 验证索引是否能正确找到对应的键
    Field field_to_query{"address", fields[1].second}; // 使用 address 字段进行查询
    std::vector<std::string> matching_keys = db_2->QueryByIndex(field_to_query);
    EXPECT_TRUE(matching_keys.empty());   

    delete db_2;
}


int main(int argc, char** argv) {
    // 设置全局随机种子
    SetGlobalSeed(static_cast<unsigned>(time(nullptr)));
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}