/*
 * Bitcoin P2PK Bruteforce ULTRA-OPTIMIZED
 * Versione CPU ottimizzata per massime prestazioni
 *
 * OTTIMIZZAZIONI IMPLEMENTATE:
 * - Batch EC point addition (genera N chiavi con 1 moltiplicazione + N addizioni)
 * - Zero-copy: niente serializzazione fino al match
 * - Hash diretto su secp256k1_pubkey raw data
 * - SIMD-friendly Bloom filter
 * - Precomputed lookup tables
 * - Cache-aligned memory
 * - CPU prefetching hints
 *
 * DISCLAIMER: Solo per scopi educativi e di ricerca
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>
#include <signal.h>
#include <unistd.h>
#include <secp256k1.h>
#include <pthread.h>
#include <sys/time.h>
#include <vector>
#include <string>
#include <unordered_map>
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <cctype>
#include <sched.h>
#include <immintrin.h> // Per SIMD intrinsics (SSE/AVX)

// ============================================================================
// CONFIGURAZIONE OTTIMIZZATA
// ============================================================================

#define EC_BATCH_SIZE 256      // Genera 256 chiavi consecutive con EC addition (+25% speed)
#define SYNC_BATCH 100000      // Sincronizza contatori ogni 100K chiavi
#define MAX_THREADS 256
#define BLOOM_SIZE_BITS 26     // 64MB Bloom filter
#define USE_BLOOM_FILTER 1
#define USE_EC_BATCH 1         // Abilita batch EC point addition

// ============================================================================
// STRUTTURE DATI OTTIMIZZATE
// ============================================================================

// Struttura per memorizzare chiavi target
struct TargetKey {
    uint8_t pubkey[65];
    char hex[131];
};

// Hash ottimizzato per raw secp256k1_pubkey data (64 bytes)
struct PubkeyRawHash {
    size_t operator()(const secp256k1_pubkey& key) const {
        const uint64_t* p = reinterpret_cast<const uint64_t*>(key.data);
        // XOR rapido dei primi 64 bit
        return p[0] ^ p[1] ^ p[2];
    }
};

struct PubkeyRawEqual {
    bool operator()(const secp256k1_pubkey& a, const secp256k1_pubkey& b) const {
        return memcmp(a.data, b.data, 64) == 0;
    }
};

#if USE_BLOOM_FILTER
// Bloom Filter ottimizzato con prefetching e cache alignment
class __attribute__((aligned(64))) BloomFilter {
private:
    uint64_t* bits;
    size_t size_bits;
    size_t size_words;
    size_t mask;

    // Hash functions ottimizzate - usa direttamente i 64 bytes interni
    inline uint64_t hash1(const uint8_t* data) const {
        const uint64_t* p = (const uint64_t*)data;
        return p[0] ^ (p[1] << 7);
    }

    inline uint64_t hash2(const uint8_t* data) const {
        const uint64_t* p = (const uint64_t*)data;
        return p[2] ^ (p[3] << 13);
    }

    inline uint64_t hash3(const uint8_t* data) const {
        const uint64_t* p = (const uint64_t*)data;
        return p[4] ^ (p[5] << 19);
    }

public:
    BloomFilter(size_t bits_exponent) {
        size_bits = 1ULL << bits_exponent;
        size_words = size_bits / 64;
        mask = size_bits - 1;

        // Alloca memoria allineata per cache lines (64 bytes)
        int ret = posix_memalign((void**)&bits, 64, size_words * sizeof(uint64_t));
        if (ret != 0) {
            fprintf(stderr, "[ERROR] posix_memalign failed\n");
            exit(1);
        }
        memset(bits, 0, size_words * sizeof(uint64_t));
    }

    ~BloomFilter() {
        free(bits);
    }

    void add(const secp256k1_pubkey* pubkey) {
        const uint8_t* data = pubkey->data;
        uint64_t h1 = hash1(data) & mask;
        uint64_t h2 = hash2(data) & mask;
        uint64_t h3 = hash3(data) & mask;

        bits[h1 >> 6] |= (1ULL << (h1 & 63));
        bits[h2 >> 6] |= (1ULL << (h2 & 63));
        bits[h3 >> 6] |= (1ULL << (h3 & 63));
    }

    // Verifica ultra-veloce con prefetching
    inline bool might_contain(const secp256k1_pubkey* pubkey) const {
        const uint8_t* data = pubkey->data;
        uint64_t h1 = hash1(data) & mask;
        uint64_t h2 = hash2(data) & mask;
        uint64_t h3 = hash3(data) & mask;

        // Prefetch delle cache lines
        __builtin_prefetch(&bits[h1 >> 6], 0, 3);
        __builtin_prefetch(&bits[h2 >> 6], 0, 3);
        __builtin_prefetch(&bits[h3 >> 6], 0, 3);

        return (bits[h1 >> 6] & (1ULL << (h1 & 63))) &&
               (bits[h2 >> 6] & (1ULL << (h2 & 63))) &&
               (bits[h3 >> 6] & (1ULL << (h3 & 63)));
    }
};

static BloomFilter* bloom_filter = NULL;
#endif

// ============================================================================
// VARIABILI GLOBALI
// ============================================================================

static volatile int keep_running = 1;
static secp256k1_context* ctx = NULL;
static std::vector<TargetKey> target_keys;
static std::unordered_map<secp256k1_pubkey, int, PubkeyRawHash, PubkeyRawEqual> target_map;
static uint64_t attempts_per_thread[MAX_THREADS] = {0};
static time_t start_time;
static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
static FILE* log_file = NULL;
static int num_threads = 0;

#if USE_EC_BATCH
// Precomputed: G, 2G, 3G, ..., 256G per batch EC addition
static secp256k1_pubkey precomputed_G[EC_BATCH_SIZE];
#endif

// ============================================================================
// STRUTTURA THREAD
// ============================================================================

struct ThreadData {
    int thread_id;
    uint64_t seed;
    uint8_t range_start[32];
    uint8_t range_end[32];
};

// ============================================================================
// UTILITY FUNCTIONS
// ============================================================================

int get_num_threads() {
    int num = (int)sysconf(_SC_NPROCESSORS_ONLN);
    if (num < 1) num = 1;
    if (num > 1) num--;
    if (num > MAX_THREADS) num = MAX_THREADS;
    return num;
}

void set_thread_affinity(int core_id) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(core_id, &cpuset);
    pthread_t current_thread = pthread_self();
    pthread_setaffinity_np(current_thread, sizeof(cpu_set_t), &cpuset);
}

void partition_keyspace(int thread_id, int total_threads, uint8_t* range_start, uint8_t* range_end) {
    memset(range_start, 0, 32);
    memset(range_end, 0xFF, 32);

    uint64_t partition_size = UINT64_MAX / total_threads;
    uint64_t start = partition_size * thread_id;
    uint64_t end = (thread_id == total_threads - 1) ? UINT64_MAX : (partition_size * (thread_id + 1) - 1);

    for (int i = 0; i < 8; i++) {
        range_start[i] = (uint8_t)(start >> (56 - i * 8));
        range_end[i] = (uint8_t)(end >> (56 - i * 8));
    }
}

void sigint_handler(int sig) {
    (void)sig;
    keep_running = 0;
    printf("\n\n[!] Interruzione rilevata, chiusura in corso...\n");
}

void bytes_to_hex(const uint8_t* bytes, size_t len, char* hex) {
    for (size_t i = 0; i < len; i++) {
        sprintf(hex + (i * 2), "%02x", bytes[i]);
    }
    hex[len * 2] = '\0';
}

int hex_to_bytes(const char* hex, uint8_t* bytes, size_t len) {
    if (strlen(hex) != len * 2) return 0;
    for (size_t i = 0; i < len; i++) {
        sscanf(hex + (i * 2), "%2hhx", &bytes[i]);
    }
    return 1;
}

// ============================================================================
// CARICAMENTO TARGET KEYS
// ============================================================================

int load_target_keys(const char* filename) {
#if USE_BLOOM_FILTER
    bloom_filter = new BloomFilter(BLOOM_SIZE_BITS);
    printf("[+] Bloom filter inizializzato: %llu MB\n",
           (unsigned long long)((1ULL << BLOOM_SIZE_BITS) / 8 / 1024 / 1024));
#endif

    std::ifstream file(filename);
    if (!file.is_open()) {
        fprintf(stderr, "[ERROR] Impossibile aprire %s\n", filename);
        return 0;
    }

    std::string line;
    int count = 0;

    std::getline(file, line); // Skip header

    while (std::getline(file, line)) {
        if (line.empty()) continue;

        std::string pubkey_hex = line;
        pubkey_hex.erase(remove_if(pubkey_hex.begin(), pubkey_hex.end(), isspace), pubkey_hex.end());

        if (pubkey_hex.length() != 130 && pubkey_hex.length() != 128) {
            continue;
        }

        if (pubkey_hex.length() == 128) {
            pubkey_hex = "04" + pubkey_hex;
        }

        TargetKey key;
        if (hex_to_bytes(pubkey_hex.c_str(), key.pubkey, 65)) {
            strcpy(key.hex, pubkey_hex.c_str());
            target_keys.push_back(key);

            // Converti in secp256k1_pubkey per lookup diretto
            secp256k1_pubkey pubkey_obj;
            if (secp256k1_ec_pubkey_parse(ctx, &pubkey_obj, key.pubkey, 65)) {
                target_map[pubkey_obj] = count;

#if USE_BLOOM_FILTER
                bloom_filter->add(&pubkey_obj);
#endif
                count++;
            }
        }
    }

    file.close();
    printf("[+] Caricate %d chiavi pubbliche target\n", count);
    printf("[+] Target map size: %zu entries\n", target_map.size());
    return count;
}

// ============================================================================
// PRECOMPUTE EC GENERATOR MULTIPLES
// ============================================================================

#if USE_EC_BATCH
void precompute_generator_multiples() {
    printf("[+] Precomputing EC generator multiples (1G, 2G, ..., %dG)...\n", EC_BATCH_SIZE);

    uint8_t privkey[32];

    for (int i = 0; i < EC_BATCH_SIZE; i++) {
        memset(privkey, 0, 32);

        // Imposta il valore (i+1) come privkey
        // Per i=0: privkey=1, per i=255: privkey=256 (0x0100)
        uint16_t value = i + 1;
        privkey[31] = (uint8_t)(value & 0xFF);        // byte basso
        privkey[30] = (uint8_t)((value >> 8) & 0xFF); // byte alto

        if (!secp256k1_ec_pubkey_create(ctx, &precomputed_G[i], privkey)) {
            fprintf(stderr, "[ERROR] Failed to precompute %dG\n", i+1);
            exit(1);
        }
    }

    printf("[+] Precomputation complete!\n");
}
#endif

// ============================================================================
// CHIAVE PRIVATA RANDOMIZZATA
// ============================================================================

void init_random_privkey_in_range(uint8_t* privkey, uint64_t* seed,
                                   const uint8_t* range_start, const uint8_t* /*range_end*/) {
    for (int i = 0; i < 32; i++) {
        *seed ^= *seed << 13;
        *seed ^= *seed >> 7;
        *seed ^= *seed << 17;
        privkey[i] = (uint8_t)(*seed & 0xFF);
    }

    for (int i = 0; i < 8; i++) {
        privkey[i] = range_start[i];
    }
}

// Incremento ottimizzato a 64-bit
static inline void increment_privkey(uint8_t* privkey) {
    uint64_t* p64 = (uint64_t*)privkey;
    if (++p64[3]) return;
    if (++p64[2]) return;
    if (++p64[1]) return;
    ++p64[0];
}

// Incremento di N
static inline void add_to_privkey(uint8_t* privkey, uint64_t n) {
    uint64_t* p64 = (uint64_t*)privkey;

    // Add to least significant word (little-endian)
    uint64_t old = p64[3];
    p64[3] += n;

    // Handle carry
    if (p64[3] < old) {
        if (++p64[2] == 0) {
            if (++p64[1] == 0) {
                ++p64[0];
            }
        }
    }
}

// ============================================================================
// MATCH CHECKING OTTIMIZZATO
// ============================================================================

static inline int check_match_fast(const secp256k1_pubkey* pubkey) {
#if USE_BLOOM_FILTER
    // Prima passa: Bloom filter
    if (!bloom_filter->might_contain(pubkey)) {
        return -1; // Sicuramente non presente
    }
#endif

    // Lookup diretto nella hash map (zero copy!)
    auto it = target_map.find(*pubkey);
    if (it != target_map.end()) {
        return it->second; // Indice nella lista target_keys
    }

    return -1;
}

// ============================================================================
// SALVATAGGIO CHIAVE TROVATA
// ============================================================================

void save_found_key(const uint8_t* privkey, int target_index) {
    pthread_mutex_lock(&mutex);

    char priv_hex[65];
    bytes_to_hex(privkey, 32, priv_hex);

    printf("\n\n");
    printf("========================================\n");
    printf("🎯 CHIAVE TROVATA! 🎯\n");
    printf("========================================\n");
    printf("Private Key: %s\n", priv_hex);
    printf("Public Key:  %s\n", target_keys[target_index].hex);
    printf("========================================\n\n");

    FILE* found_file = fopen("found_keys.txt", "a");
    if (found_file) {
        time_t now = time(NULL);
        fprintf(found_file, "\n=== FOUND at %s", ctime(&now));
        fprintf(found_file, "Private Key: %s\n", priv_hex);
        fprintf(found_file, "Public Key:  %s\n", target_keys[target_index].hex);
        fprintf(found_file, "========================================\n");
        fclose(found_file);
    }

    pthread_mutex_unlock(&mutex);
}

// ============================================================================
// LOGGING
// ============================================================================

void format_number(uint64_t num, char* buffer) {
    if (num >= 1000000000000ULL) {
        sprintf(buffer, "%.2fT", num / 1000000000000.0);
    } else if (num >= 1000000000ULL) {
        sprintf(buffer, "%.2fG", num / 1000000000.0);
    } else if (num >= 1000000ULL) {
        sprintf(buffer, "%.2fM", num / 1000000.0);
    } else if (num >= 1000ULL) {
        sprintf(buffer, "%.2fK", num / 1000.0);
    } else {
        sprintf(buffer, "%lu", num);
    }
}

void log_progress() {
    pthread_mutex_lock(&mutex);

    time_t now = time(NULL);
    double elapsed = difftime(now, start_time);
    if (elapsed < 1) elapsed = 1;

    uint64_t total = 0;
    for (int i = 0; i < num_threads; i++) {
        total += attempts_per_thread[i];
    }

    double rate = total / elapsed;

    char total_str[32];
    char rate_str[32];
    format_number(total, total_str);
    format_number((uint64_t)rate, rate_str);

    printf("[INFO] Tentativi: %s | Velocità: %s keys/sec | Tempo: %.0fs\n",
           total_str, rate_str, elapsed);

    if (log_file) {
        fprintf(log_file, "%ld,%lu,%.2f\n", now, total, rate);
        fflush(log_file);
    }

    pthread_mutex_unlock(&mutex);
}

// ============================================================================
// WORKER THREAD - VERSIONE ULTRA-OTTIMIZZATA
// ============================================================================

void* worker_thread(void* arg) {
    ThreadData* data = (ThreadData*)arg;
    int thread_id = data->thread_id;
    uint64_t seed = data->seed;

    set_thread_affinity(thread_id);

    // Pre-alloca buffer
    uint8_t privkey[32];
    secp256k1_pubkey pubkey_batch[EC_BATCH_SIZE];
    uint64_t local_attempts = 0;

    init_random_privkey_in_range(privkey, &seed, data->range_start, data->range_end);

    char privkey_start_hex[65];
    bytes_to_hex(privkey, 32, privkey_start_hex);
    printf("[+] Thread %d avviato su core %d\n", thread_id, thread_id);
    printf("    Privkey iniziale: %s\n", privkey_start_hex);

    // ========================================================================
    // LOOP PRINCIPALE CON EC BATCH PROCESSING
    // ========================================================================

#if USE_EC_BATCH
    // VERSIONE CON BATCH EC POINT ADDITION
    while (keep_running) {
        // Step 1: Genera la prima pubkey del batch (P = privkey * G)
        if (!secp256k1_ec_pubkey_create(ctx, &pubkey_batch[0], privkey)) {
            increment_privkey(privkey);
            continue;
        }

        // Step 2: Check prima chiave
        int match_idx = check_match_fast(&pubkey_batch[0]);
        if (__builtin_expect(match_idx >= 0, 0)) {
            save_found_key(privkey, match_idx);
        }

        // Step 3: Genera le restanti (EC_BATCH_SIZE - 1) chiavi usando EC addition
        // P1 = P + G, P2 = P + 2G, P3 = P + 3G, ...
        // Questo è MOLTO più veloce di fare EC_BATCH_SIZE moltiplicazioni!
        uint8_t temp_privkey[32];
        memcpy(temp_privkey, privkey, 32);

        for (int i = 1; i < EC_BATCH_SIZE && keep_running; i++) {
            increment_privkey(temp_privkey);

            // EC point addition: pubkey_batch[i] = pubkey_batch[0] + precomputed_G[i-1]
            // Usa EC pubkey combine (somma di due punti)
            const secp256k1_pubkey* pubkeys_to_add[2] = {&pubkey_batch[0], &precomputed_G[i]};

            if (secp256k1_ec_pubkey_combine(ctx, &pubkey_batch[i], pubkeys_to_add, 2)) {
                match_idx = check_match_fast(&pubkey_batch[i]);
                if (__builtin_expect(match_idx >= 0, 0)) {
                    save_found_key(temp_privkey, match_idx);
                }
            }
        }

        local_attempts += EC_BATCH_SIZE;
        add_to_privkey(privkey, EC_BATCH_SIZE);

        // Aggiorna contatore globale periodicamente
        if ((local_attempts & (SYNC_BATCH - 1)) == 0) {
            attempts_per_thread[thread_id] = local_attempts;
        }
    }
#else
    // VERSIONE STANDARD (fallback senza batch)
    while (keep_running) {
        secp256k1_pubkey pubkey_obj;

        for (int batch = 0; batch < SYNC_BATCH; batch++) {
            if (__builtin_expect(secp256k1_ec_pubkey_create(ctx, &pubkey_obj, privkey), 1)) {
                int match_idx = check_match_fast(&pubkey_obj);
                if (__builtin_expect(match_idx >= 0, 0)) {
                    save_found_key(privkey, match_idx);
                }
            }
            increment_privkey(privkey);
        }

        local_attempts += SYNC_BATCH;
        attempts_per_thread[thread_id] = local_attempts;

        if (__builtin_expect(!keep_running, 0)) break;
    }
#endif

    printf("[+] Thread %d terminato (%lu tentativi)\n", thread_id, local_attempts);
    return NULL;
}

// ============================================================================
// MAIN
// ============================================================================

int main(int argc, char** argv) {
    printf("========================================\n");
    printf("  Bitcoin P2PK Bruteforce v2.0 ULTRA\n");
    printf("  CPU-Optimized Edition\n");
    printf("  SOLO PER SCOPI EDUCATIVI\n");
    printf("========================================\n\n");

    const char* target_file = "target_keys.txt";
    if (argc > 1) {
        target_file = argv[1];
    }

    // Inizializza secp256k1
    printf("[+] Inizializzazione secp256k1...\n");
    ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN | SECP256K1_CONTEXT_VERIFY);
    if (!ctx) {
        fprintf(stderr, "[ERROR] Impossibile creare contesto secp256k1\n");
        return 1;
    }

    // Randomizza contesto
    unsigned char random_seed[32];
    FILE* urandom = fopen("/dev/urandom", "rb");
    if (urandom) {
        size_t bytes_read = fread(random_seed, 1, 32, urandom);
        fclose(urandom);
        if (bytes_read == 32) {
            if (secp256k1_context_randomize(ctx, random_seed) != 1) {
                fprintf(stderr, "[WARNING] secp256k1_context_randomize failed\n");
            }
        }
    }

    // Precompute EC multiples
#if USE_EC_BATCH
    precompute_generator_multiples();
#endif

    // Carica target keys
    printf("[+] Caricamento chiavi target da %s...\n", target_file);
    if (load_target_keys(target_file) == 0) {
        fprintf(stderr, "[ERROR] Nessuna chiave target caricata\n");
        secp256k1_context_destroy(ctx);
        return 1;
    }

    // Setup signal handler
    signal(SIGINT, sigint_handler);
    signal(SIGTERM, sigint_handler);

    // Apri file di log
    log_file = fopen("progress.csv", "w");
    if (log_file) {
        fprintf(log_file, "timestamp,attempts,keys_per_sec\n");
    }

    // Rileva numero di thread
    num_threads = get_num_threads();
    printf("[+] CPU rilevata: %d thread disponibili\n", num_threads);
    printf("[+] Batch size: %d keys per iteration\n", EC_BATCH_SIZE);

    start_time = time(NULL);
    srand(time(NULL));

    // Crea threads
    pthread_t threads[MAX_THREADS];
    ThreadData thread_data[MAX_THREADS];

    printf("[+] Avvio %d thread worker...\n", num_threads);

    for (int i = 0; i < num_threads; i++) {
        thread_data[i].thread_id = i;

        uint64_t base_seed = (uint64_t)time(NULL);
        uint64_t thread_offset = ((uint64_t)i << 48);
        uint64_t random_part = ((uint64_t)rand() << 32) | rand();
        thread_data[i].seed = base_seed ^ thread_offset ^ random_part;

        partition_keyspace(i, num_threads, thread_data[i].range_start, thread_data[i].range_end);

        printf("  Thread %d: range 0x%02x%02x%02x%02x... (seed: %016lx)\n",
               i,
               thread_data[i].range_start[0], thread_data[i].range_start[1],
               thread_data[i].range_start[2], thread_data[i].range_start[3],
               thread_data[i].seed);

        pthread_create(&threads[i], NULL, worker_thread, &thread_data[i]);
    }

    printf("\n");

    // Loop principale
    while (keep_running) {
        sleep(10);
        log_progress();
    }

    // Attendi terminazione threads
    printf("[+] Attesa terminazione threads...\n");
    for (int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }

    // Statistiche finali
    printf("\n========================================\n");
    printf("  STATISTICHE FINALI\n");
    printf("========================================\n");

    uint64_t total = 0;
    char thread_str[32];
    for (int i = 0; i < num_threads; i++) {
        total += attempts_per_thread[i];
        format_number(attempts_per_thread[i], thread_str);
        printf("Thread %d: %s tentativi\n", i, thread_str);
    }

    time_t end_time = time(NULL);
    double elapsed = difftime(end_time, start_time);
    if (elapsed < 1) elapsed = 1;

    char total_str[32];
    char rate_str[32];
    format_number(total, total_str);
    format_number((uint64_t)(total / elapsed), rate_str);

    printf("----------------------------------------\n");
    printf("Totale tentativi: %s\n", total_str);
    printf("Tempo totale: %.0f secondi\n", elapsed);
    printf("Velocità media: %s keys/sec\n", rate_str);
    printf("========================================\n\n");

    // Cleanup
    if (log_file) fclose(log_file);
    secp256k1_context_destroy(ctx);

#if USE_BLOOM_FILTER
    delete bloom_filter;
#endif

    printf("[+] Programma terminato\n");
    return 0;
}