#include <cassert>
#include <iostream>
#include <map>
#include <thread>
#include <random>
#include <chrono>
#include <atomic>
#include "UnorderedParallelMap.h"

constexpr int THREADS = 8;
constexpr int OPS_PER_THREAD = 10000;

std::atomic<uint64_t> insert_ns{0};
std::atomic<uint64_t> find_ns{0};
std::atomic<uint64_t> update_ns{0};
std::atomic<uint64_t> erase_ns{0};

inline uint64_t now_ns() {
    return std::chrono::duration_cast<std::chrono::nanoseconds>(
        std::chrono::steady_clock::now().time_since_epoch()).count();
}

void worker(LockFreeMap<int, int>& map, int thread_id)
{
    std::mt19937 rng(thread_id);
    std::uniform_int_distribution<int> dist(0, 99999);

    for (int i = 0; i < OPS_PER_THREAD; ++i)
    {
        int op = dist(rng) % 4;
        int key = dist(rng);
        int value = thread_id * 1000 + i;

        uint64_t start = now_ns();

        switch (op)
        {
        case 0:
            map.insert(key, value);
            insert_ns += (now_ns() - start);
            break;
        case 1:
            map.find(key);
            find_ns += (now_ns() - start);
            break;
        case 2:
            map.update(key, value + 1);
            update_ns += (now_ns() - start);
            break;
        case 3:
            map.erase(key);
            erase_ns += (now_ns() - start);
            break;
        }
    }
}

void verify_map(LockFreeMap<int, int>& map)
{
    std::cout << "Final size: " << map.size() << "\n";

    auto s = map.size();
    auto k = map.keys().size();
    auto e = map.entries().size();

    std::cout << "Verifying invariants...\n";
    std::cout << " size() = " << s << "\n";
    std::cout << " keys() = " << k << "\n";
    std::cout << " entries() = " << e << "\n";

    assert(e <= k);
    assert(k <= s);
    std::cout << "✓ Invariants OK\n\n";
}

int main()
{
    LockFreeMap<int, int> simple_test_map;
    simple_test_map.insert(1, 1);
    simple_test_map.insert(2, 44);
    assert(simple_test_map.size() == 2);
    assert(simple_test_map.find(2).value_or(-1) == 44);
    simple_test_map.update(2, 55);
    assert(simple_test_map.find(2).value_or(-1) == 55);
    std::cout << "✓ Single-thread test finished\n" << std::endl;
    verify_map(simple_test_map);

    LockFreeMap<int, int> map(1024);
    std::vector<std::thread> threads;
    threads.reserve(THREADS);

    auto start_total = std::chrono::steady_clock::now();

    for (int i = 0; i < THREADS; ++i)
        threads.emplace_back(worker, std::ref(map), i);

    for (auto& t : threads)
        t.join();

    auto end_total = std::chrono::steady_clock::now();
    auto total_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(end_total - start_total).count();

    std::cout << "✓ Multi-thread test finished\n" << std::endl;

    verify_map(map);

    uint64_t total_ops = THREADS * OPS_PER_THREAD;
    std::cout << "--- Benchmark Results ---\n";
    std::cout << "Total ops:      " << total_ops << "\n";
    std::cout << "Total time:     " << total_ns / 1e6 << " ms\n";
    std::cout << "Insert avg:     " << insert_ns / 1e3 << " μs total (" << (insert_ns / (double)total_ops) << " ns/op)\n";
    std::cout << "Find avg:       " << find_ns / 1e3 << " μs total (" << (find_ns / (double)total_ops) << " ns/op)\n";
    std::cout << "Update avg:     " << update_ns / 1e3 << " μs total (" << (update_ns / (double)total_ops) << " ns/op)\n";
    std::cout << "Erase avg:      " << erase_ns / 1e3 << " μs total (" << (erase_ns / (double)total_ops) << " ns/op)\n";
    std::cout << "--------------------------\n";

    return 0;
}