#include <iostream>
#include <vector>
#include <string>
#include <chrono>
#include <unordered_map>
#include <iomanip>

#define HM_IMPLEMENTATION
#include "hm.h"

using namespace std::chrono;

typedef std::chrono::time_point<std::chrono::high_resolution_clock> time_p;

// Configuration
const int NUM_ELEMENTS = 1000000;

time_p now()
{
    return high_resolution_clock::now();
}

double time_diff_nano(time_p end, time_p start)
{
    return duration_cast<nanoseconds>(end - start).count() / 1e6;
}

void print_div()
{
    std::cout << "---------------------------------------------------------" << std::endl;
}

void print_head()
{
    std::cout << std::left << std::setw(10) << "Operation"
        << std::right << std::setw(22) << "std::unordered_map"
        << std::right << std::setw(13) << "hm.h"
        << std::right << std::setw(11) << "Ratio" << std::endl;
}

void print_row(const std::string& label, double hm_ms, double std_ms)
{
    std::cout << std::left << std::setw(10) << label
        << std::right << std::setw(19) << std::fixed << std::setprecision(2) << std_ms << " ms"
        << std::right << std::setw(10) << std::fixed << std::setprecision(2) << hm_ms << " ms"
        << std::right << std::setw(10) << (hm_ms / std_ms) << "x" << std::endl;
}

// Test with string key type
uint64_t hash_str(const void *buffer, size_t size)
{
    (void)size;
    const char *str = reinterpret_cast<const char *>(buffer);
    uint64_t hash = (uint64_t)str;
    return hash;
}

void test_str()
{
    std::vector<std::string> key_pool;
    std::vector<std::string> value_pool;

    // Generate random strings to ensure we aren't just hashing "1", "2", "3"
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        key_pool.push_back("key_prefix_" + std::to_string(i) + "_suffix_" + std::to_string(rand() % 1000));
        value_pool.push_back("value_" + std::to_string(i));
    }

    // --- Custom implementation ---
    // Using 0 for sizes tells your library to treat them as NULL-terminated strings
    HashMap hm = hm_create_ex(sizeof(const char *), 0, hash_str);

    // Insertion
    auto start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        hm_put(&hm, key_pool[i].c_str(), value_pool[i].c_str());
    }
    auto end = now();
    double custom_put = time_diff_nano(end, start);

    // Lookup
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        char* val = (char*)hm_get(&hm, key_pool[i].c_str());
        if (!val) {
            std::cerr << "Lookup error!" << std::endl;
            return;
        }
    }
    end = now();
    double custom_get = time_diff_nano(end, start);

    // Deletion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        hm_remove(&hm, key_pool[i].c_str());
    }
    end = now();
    double custom_del = time_diff_nano(end, start);

    hm_free(&hm);

    // --- Benchmark std::unordered_map ---
    std::unordered_map<const char*, const char*> std_map;

    // Insertion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        std_map[key_pool[i].c_str()] = value_pool[i].c_str();
    }
    end = now();
    double std_put = time_diff_nano(end, start);

    // Lookup
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        auto it = std_map.find(key_pool[i].c_str());
        if (it == std_map.end()) std::cerr << "Lookup error!" << std::endl;
    }
    end = now();
    double std_get = time_diff_nano(end, start);

    // Deletion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        std_map.erase(key_pool[i].c_str());
    }
    end = now();
    double std_del = time_diff_nano(end, start);

    // Output Results
    print_head();
    print_div();
    print_row("Insert", custom_put, std_put);
    print_row("Lookup", custom_get, std_get);
    print_row("Delete", custom_del, std_del);
    print_div();
}

// Test with integer key type
uint64_t hash_int(const void *buffer, size_t size)
{
    (void)size;
    return *reinterpret_cast<const int*>(buffer);
}

void test_int()
{
    // Prepare test data (random integers) to avoid string allocation overhead in the loops
    std::vector<int> keys(NUM_ELEMENTS);
    std::vector<int> values(NUM_ELEMENTS);

    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        keys[i] = i;
        values[i] = i * 2;
    }

    // --- Custom implementation ---
    HashMap hm = hm_create_ex(sizeof(int), sizeof(int), hash_int);

    // Insertion
    auto start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        hm_put(&hm, &keys[i], &values[i]);
    }
    auto end = now();
    double custom_put = time_diff_nano(end, start);

    // Lookup
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        int* val = (int*)hm_get(&hm, &keys[i]);
        if (!val || *val != values[i]) {
            std::cerr << "Error in custom lookup!" << std::endl;
            return;
        }
    }
    end = now();
    double custom_get = time_diff_nano(end, start);

    // Deletion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        hm_remove(&hm, &keys[i]);
    }
    end = now();
    double custom_del = time_diff_nano(end, start);

    hm_free(&hm);

    // --- Benchmark std::unordered_map ---
    std::unordered_map<int, int> std_map;

    // Insertion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        std_map[keys[i]] = values[i];
    }
    end = now();
    double std_put = time_diff_nano(end, start);

    // Lookup
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        auto it = std_map.find(keys[i]);
        if (it == std_map.end() || it->second != values[i]) {
            std::cerr << "Error in std lookup!" << std::endl;
        }
    }
    end = now();
    double std_get = time_diff_nano(end, start);

    // Deletion
    start = now();
    for (int i = 0; i < NUM_ELEMENTS; ++i) {
        std_map.erase(keys[i]);
    }
    end = now();
    double std_del = time_diff_nano(end, start);

    // Output Results
    print_head();
    print_div();
    print_row("Insert", custom_put, std_put);
    print_row("Lookup", custom_get, std_get);
    print_row("Delete", custom_del, std_del);
    print_div();
}

int main()
{
    std::cout << "Benchmarking " << NUM_ELEMENTS << " (integer keys)" << std::endl << std::endl;
    test_int();

    std::cout << std::endl;

    std::cout << "Benchmarking " << NUM_ELEMENTS << " (string keys)" << std::endl << std::endl;
    test_str();

    return 0;
}