#include <iostream>
#include <fstream>
#include <ctime>
#include <chrono>
#include <filesystem>
#include <vector>
#include <thread>
#include <cmath>
#include <mutex>
#include <atomic>
#include "Ec.h"
#include "bloom/filter.hpp"

using namespace std;
namespace fs = filesystem;

mutex cout_mutex;
atomic<bool> stop_flag(false);

inline uint64_t str_to_uint64(const string& value) {
    try {
        return stoull(value);
    } catch (...) {
        cerr << "Invalid number format" << endl;
        exit(1);
    }
}

string trim(const string& str) {
    auto start = str.begin();
    while (start != str.end() && isspace(*start)) ++start;
    auto end = str.end();
    do { --end; } while (end != start && isspace(*end));
    return string(start, end + 1);
}

void print_time() {
    time_t timestamp = time(NULL);
    struct tm datetime = *localtime(&timestamp);
    char output[50];
    strftime(output, 50, "%H:%M:%S", &datetime);
    cout << "[" << output << "] ";
}

void bloom_creator(int index, EcPoint start_point, uint64_t n_elements, double error, const string& strideS) {
    Ec ec; InitEc();
    string filename = "bloom" + to_string(index) + ".bf";
    string settingsfile = "settings" + to_string(index) + ".txt";

    try {
        using filter = boost::bloom::filter<string, 32>;
        filter bf(n_elements, error);
        EcPoint P = start_point;
        extern EcPoint g_G;

        {
            lock_guard<mutex> lock(cout_mutex);
            print_time(); cout << "Creating " << filename << endl;
        }

        for (uint64_t i = 0; i < n_elements && !stop_flag; ++i) {
            bf.insert(ec.GetPublicKeyHex(P));
            P = ec.AddPoints(P, g_G);

            if (i % 1000000 == 0) {
                lock_guard<mutex> lock(cout_mutex);
                print_time(); cout << filename << " progress: " << i << "/" << n_elements << endl;
            }
        }

        if (stop_flag) return;

        ofstream out(filename, ios::binary);
        size_t cap = bf.capacity();
        out.write(reinterpret_cast<const char*>(&cap), sizeof(cap));
        boost::span<const unsigned char> s = bf.array();
        out.write(reinterpret_cast<const char*>(s.data()), s.size());
        out.close();

        {
            lock_guard<mutex> lock(cout_mutex);
            print_time(); cout << "Finished writing " << filename << endl;
        }

        // Save settings (matching 2-way version)
        ofstream settings_out(settingsfile);
        settings_out << ec.GetPublicKeyHex(start_point) << '\n';
        settings_out << strideS << '\n';
        settings_out.close();

    } catch (const exception& e) {
        lock_guard<mutex> lock(cout_mutex);
        cerr << "Error in thread " << index << ": " << e.what() << endl;
        stop_flag = true;
    }
}

int main() {
    auto start = chrono::high_resolution_clock::now();
    Ec ec; InitEc();

    // Clean up old files
    for (int i = 0; i < 8; ++i) {
        remove(("bloom" + to_string(i) + ".bf").c_str());
        remove(("settings" + to_string(i) + ".txt").c_str());
    }

    // Generate P_table (like 2-way version)
    vector<EcPoint> P_table;
    EcInt pk; pk.Set(1);
    uint64_t mult = 2;
    EcPoint P;
    for (int i = 0; i < 256; i++) {
        P = ec.MultiplyG(pk);
        P_table.push_back(P);
        pk.Mul_u64(pk, mult);
    }

    uint64_t range_start, range_end, block_width;
    string temp, search_pub;
    ifstream inFile("settings.txt");
    if (!inFile) {
        cerr << "Cannot open settings.txt" << endl;
        return 1;
    }

    getline(inFile, temp); range_start = str_to_uint64(temp);
    getline(inFile, temp); range_end = str_to_uint64(temp);
    getline(inFile, temp); block_width = str_to_uint64(temp);
    getline(inFile, temp); search_pub = trim(temp);
    inFile.close();

    uint64_t n_elements = uint64_t(pow(2, block_width));
    double error = 1e-8;

    // Calculate puzzle_point_divide2 (like 2-way version)
    EcPoint puzzle_point;
    puzzle_point.SetHexStr(search_pub.c_str());
    EcInt d_05;
    d_05.SetHexStr("7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF5D576E7357A4501DDFE92F46681B20A1");
    EcPoint puzzle_point_divide2 = ec.MultiplyP(puzzle_point, d_05);

    // Calculate Q2 (like 2-way version)
    EcPoint first_point = P_table[range_start - 1];
    EcPoint second_point = P_table[range_start - 2];
    EcPoint P1 = ec.SubtractPoints(puzzle_point_divide2, first_point);
    EcPoint P2 = ec.SubtractPoints(puzzle_point_divide2, second_point);
    EcPoint Q1 = ec.AddPoints(P1, P2);
    EcPoint Q2 = ec.AddPoints(puzzle_point_divide2, Q1);

    // Prepare stride (like 2-way version)
    EcInt stride_sum; stride_sum.Set(0);
    char strideS[100];
    stride_sum.GetHexStr(strideS);

    // Generate starting points (4 forward, 4 backward from Q2)
    vector<EcPoint> start_points;
    EcInt stride; stride.Set(n_elements);
    EcPoint strideG = ec.MultiplyG(stride);

    // Forward points (0-3)
    EcPoint base = Q2;
    for (int i = 0; i < 4; ++i) {
        start_points.push_back(base);
        base = ec.AddPoints(base, strideG);
    }

    // Backward points (4-7)
    base = Q2;
    for (int i = 0; i < 4; ++i) {
        base = ec.SubtractPoints(base, strideG);
        start_points.push_back(base);
    }

    vector<thread> threads;
    for (int i = 0; i < 8; ++i) {
        threads.emplace_back(bloom_creator, i, start_points[i], n_elements, error, strideS);
    }

    for (auto& t : threads) {
        if (t.joinable()) t.join();
    }

    auto end = chrono::high_resolution_clock::now();
    auto duration = chrono::duration_cast<chrono::seconds>(end - start);
    print_time(); cout << "Total time: " << duration.count() << " seconds" << endl;
    return 0;
}