#include "gtest/gtest.h"
#include "db/NewDB.h"  // NewDB 的头文件
#include "leveldb/env.h"
#include "leveldb/db.h"
#include <chrono>  // 用于计时
#include <random>  // 用于生成随机数
#include <vector>  // 用于存储延迟

using namespace std;
using namespace leveldb;

// 随机生成字符串的辅助函数
std::string GenerateRandomString(size_t length) {
    const char charset[] = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
    const size_t max_index = sizeof(charset) - 1;
    std::random_device rd;
    std::mt19937 generator(rd());
    std::uniform_int_distribution<size_t> distribution(0, max_index);

    std::string random_str;
    for (size_t i = 0; i < length; ++i) {
        random_str += charset[distribution(generator)];
    }
    return random_str;
}

// 打开数据库的辅助函数
Status OpenDB(std::string dbName, NewDB** db) {
    Options options;
    options.create_if_missing = true;
    Status s = NewDB::Open(options, dbName, db);
    if (!s.ok()) {
        cerr << "Error opening database: " << s.ToString() << endl;
    }
    return s;
}

std::string testdbname = "testdb_for_performance";

// 延迟测试
void TestLatency(int num_operations, std::vector<int64_t>& lat_res) {
    int64_t latency = 0;
    auto end_time = std::chrono::steady_clock::now();
    auto last_time = end_time;

    for (int i = 0; i < num_operations; ++i) {
        // 记录操作的延迟
        end_time = std::chrono::steady_clock::now();
        latency = std::chrono::duration_cast<std::chrono::milliseconds>(
            end_time - last_time).count();
        last_time = end_time;

        lat_res.emplace_back(latency);
    }
}

// 吞吐量测试
TEST(TestNewDB, PutFieldsThroughputTest) {
    // 创建 NewDB 实例
    NewDB* db;
    ASSERT_TRUE(OpenDB(testdbname, &db).ok());

    // 插入索引字段
    db->CreateIndexOnField("address");

    // 延迟测试
    vector<int64_t> latencies;  // 存储每次插入的延迟
    int num_inserts = 1000;  // 插入的条数

    // 测试插入吞吐量
    auto start_time = std::chrono::high_resolution_clock::now();  // 记录开始时间
    for (int i = 0; i < num_inserts; ++i) {
        std::string key = "k_" + std::to_string(i);

        // 随机生成字段内容
        FieldArray fields = {
            {"name", "Customer#" + std::to_string(i)},
            {"address", GenerateRandomString(10)},  // 随机生成地址
            {"phone", GenerateRandomString(12)}     // 随机生成电话号码
        };

        db->Put_fields(WriteOptions(), key, fields);
    }

    // 延迟测试
    TestLatency(num_inserts, latencies);

    auto end_time = std::chrono::high_resolution_clock::now();  // 记录结束时间
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
    
    // 计算吞吐量并输出
    cout << "Throughput: " << num_inserts * 1000 / duration << " OPS" << endl;

    // 输出延迟
    // for (size_t i = 0; i < latencies.size(); ++i) {
    //     cout << "Latency for operation " << i + 1 << ": " << latencies[i] << " ms" << endl;
    // }
    cout << "Latency for operation " << latencies.size()-1 << ": " << latencies[latencies.size()-1] << " ms" << endl;
    // 清理数据库
    delete db;
}

// 主函数运行所有测试
int main(int argc, char** argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}