/*
 * Bitcoin P2PK Bruteforce - Ricerca chiavi private
 * Utilizza libsecp256k1 per massima efficienza
 *
 * 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 <array>
#include <unordered_set>
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <cctype>
#include <sched.h> // Per CPU affinity

// Configurazione
#define BATCH_SIZE 100000  // Batch più grande per ridurre overhead di sincronizzazione
#define SAVE_INTERVAL 300 // Salva progresso ogni 5 minuti
#define PROGRESS_INTERVAL 1000000 // Mostra progresso ogni N tentativi
#define MAX_THREADS 256 // Massimo numero di thread supportati

// Ottimizzazioni avanzate
#define USE_BLOOM_FILTER 1  // Usa Bloom filter per lookup ultra-veloce
#define BLOOM_SIZE_BITS 26  // 2^26 = 64MB bloom filter (adattare in base alla RAM)

// Struttura per memorizzare le chiavi pubbliche target
struct TargetKey {
    uint8_t pubkey[65]; // Chiave pubblica non compressa (65 bytes)
    char hex[131];      // Rappresentazione hex
};

// Hash personalizzato per array di 65 bytes (pubkey)
struct PubkeyHash {
    size_t operator()(const std::array<uint8_t, 65>& key) const {
        // Hash veloce usando i primi 8 bytes della pubkey
        const uint64_t* p = reinterpret_cast<const uint64_t*>(key.data());
        return p[0] ^ p[1];
    }
};

#if USE_BLOOM_FILTER
// Bloom Filter ultra-veloce per ridurre lookup costosi
class BloomFilter {
private:
    uint64_t* bits;
    size_t size_bits;
    size_t size_words;

    // Hash functions ottimizzate
    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;
        bits = new uint64_t[size_words]();
    }

    ~BloomFilter() {
        delete[] bits;
    }

    void add(const uint8_t* pubkey) {
        uint64_t h1 = hash1(pubkey) & (size_bits - 1);
        uint64_t h2 = hash2(pubkey) & (size_bits - 1);
        uint64_t h3 = hash3(pubkey) & (size_bits - 1);

        bits[h1 / 64] |= (1ULL << (h1 % 64));
        bits[h2 / 64] |= (1ULL << (h2 % 64));
        bits[h3 / 64] |= (1ULL << (h3 % 64));
    }

    inline bool might_contain(const uint8_t* pubkey) const {
        uint64_t h1 = hash1(pubkey) & (size_bits - 1);
        uint64_t h2 = hash2(pubkey) & (size_bits - 1);
        uint64_t h3 = hash3(pubkey) & (size_bits - 1);

        return (bits[h1 / 64] & (1ULL << (h1 % 64))) &&
               (bits[h2 / 64] & (1ULL << (h2 % 64))) &&
               (bits[h3 / 64] & (1ULL << (h3 % 64)));
    }
};

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_set<std::array<uint8_t, 65>, PubkeyHash> target_set;
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; // Numero effettivo di thread da usare

// Struttura per i thread
struct ThreadData {
    int thread_id;
    uint64_t seed;
    uint8_t range_start[32];  // Inizio range dello spazio delle chiavi
    uint8_t range_end[32];    // Fine range dello spazio delle chiavi
};

// Rileva numero di thread/core disponibili
// Lascia un thread libero per il sistema operativo e I/O
int get_num_threads() {
    int num = (int)sysconf(_SC_NPROCESSORS_ONLN);
    if (num < 1) num = 1;
    if (num > 1) num--; // Lascia un core libero per migliorare l'efficienza
    if (num > MAX_THREADS) num = MAX_THREADS;
    return num;
}

// Imposta affinity del thread a un core specifico
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();
    if (pthread_setaffinity_np(current_thread, sizeof(cpu_set_t), &cpuset) != 0) {
        fprintf(stderr, "[WARNING] Impossibile impostare affinity per core %d\n", core_id);
    }
}

// Partiziona lo spazio delle chiavi tra i thread
void partition_keyspace(int thread_id, int total_threads, uint8_t* range_start, uint8_t* range_end) {
    // Azzera entrambi gli array
    memset(range_start, 0, 32);
    memset(range_end, 0xFF, 32);

    // Partiziona usando i primi 8 bytes (64 bit)
    // Questo dà 2^64 / total_threads chiavi per thread
    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);

    // Converti in big-endian per i primi 8 bytes
    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));
    }

    // I restanti 24 bytes rimangono:
    // range_start[8..31] = 0x00 (minimo)
    // range_end[8..31] = 0xFF (massimo)
}

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

// Converti bytes in hex
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';
}

// Converti hex in bytes
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;
}

// Carica le chiavi pubbliche P2PK dal file
int load_target_keys(const char* filename) {
#if USE_BLOOM_FILTER
    // Inizializza 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;

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

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

        // Estrai la chiave pubblica (formato: hex della pubkey)
        // Il file dovrebbe contenere una pubkey per riga
        std::string pubkey_hex = line;

        // Rimuovi spazi bianchi
        pubkey_hex.erase(remove_if(pubkey_hex.begin(), pubkey_hex.end(), isspace), pubkey_hex.end());

        // P2PK non compresso: 65 bytes (130 caratteri hex)
        // Formato: 04 + 32 bytes X + 32 bytes Y
        if (pubkey_hex.length() != 130 && pubkey_hex.length() != 128) {
            continue; // Skip se non è una pubkey valida
        }

        // Aggiungi 04 se manca (formato non compresso)
        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);

            // Inserisci nel set usando std::array per lookup veloce
            std::array<uint8_t, 65> pubkey_array;
            memcpy(pubkey_array.data(), key.pubkey, 65);
            target_set.insert(pubkey_array);

#if USE_BLOOM_FILTER
            // Aggiungi anche al Bloom filter
            bloom_filter->add(key.pubkey);
#endif
            count++;
        }
    }

    file.close();
    printf("[+] Caricate %d chiavi pubbliche target\n", count);
    return count;
}

// Inizializza una chiave privata casuale nel range assegnato al thread
void init_random_privkey_in_range(uint8_t* privkey, uint64_t* seed,
                                   const uint8_t* range_start, const uint8_t* /*range_end*/) {
    // Genera 32 bytes completamente casuali usando xorshift64
    for (int i = 0; i < 32; i++) {
        *seed ^= *seed << 13;
        *seed ^= *seed >> 7;
        *seed ^= *seed << 17;
        privkey[i] = (uint8_t)(*seed & 0xFF);
    }

    // Applica il prefisso del range ai primi 8 bytes per partizionare lo spazio
    for (int i = 0; i < 8; i++) {
        privkey[i] = range_start[i];
    }
    // I restanti 24 bytes (192 bit) sono casuali all'interno del chunk del thread
}

// Incrementa la chiave privata di 1 (big-endian a 256 bit)
// Ottimizzato per architetture a 64-bit usando operazioni native
static inline void increment_privkey(uint8_t* privkey) {
    // Converti in array di uint64_t per operazioni a 64-bit (4x più veloce)
    uint64_t* p64 = (uint64_t*)privkey;

    // Incrementa partendo dal uint64_t meno significativo (little-endian in memoria)
    // privkey[24-31] = p64[3], privkey[16-23] = p64[2], ecc.
    if (++p64[3]) return; // Nessun carry nel primo blocco (caso più comune ~99.99%)
    if (++p64[2]) return; // Carry solo nel secondo blocco
    if (++p64[1]) return; // Carry solo nel terzo blocco
    ++p64[0];             // Carry fino al quarto blocco
}

// Verifica se la pubkey corrisponde a un target
// Ultra-ottimizzato: Bloom filter first, poi verifica precisa
static inline int check_match(const uint8_t* pubkey) {
#if USE_BLOOM_FILTER
    // First pass: Bloom filter (velocissimo, O(1) con 3 operazioni bit)
    if (!bloom_filter->might_contain(pubkey)) {
        return 0; // Sicuramente non presente (99.9%+ dei casi)
    }
    // Possibile match: verifica precisa con hash set
#endif

    // Verifica precisa solo se Bloom filter dice "forse presente"
    std::array<uint8_t, 65> pubkey_array;
    memcpy(pubkey_array.data(), pubkey, 65);
    return target_set.find(pubkey_array) != target_set.end();
}

// Salva una chiave trovata
void save_found_key(const uint8_t* privkey, const uint8_t* pubkey) {
    pthread_mutex_lock(&mutex);

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

    // Stampa a schermo
    printf("\n\n");
    printf("========================================\n");
    printf("🎯 CHIAVE TROVATA! 🎯\n");
    printf("========================================\n");
    printf("Private Key: %s\n", priv_hex);
    printf("Public Key:  %s\n", pub_hex);
    printf("========================================\n\n");

    // Salva su file
    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", pub_hex);
        fprintf(found_file, "========================================\n");
        fclose(found_file);
    }

    pthread_mutex_unlock(&mutex);
}

// Formatta numero con suffisso K, M, G, T
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);
    }
}

// Log progresso
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);
}

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

    // Fissa questo thread a un core specifico per massima efficienza
    set_thread_affinity(thread_id);

    // Pre-alloca tutte le variabili per evitare allocazioni nel loop
    uint8_t privkey[32];
    uint8_t pubkey[65];
    secp256k1_pubkey pubkey_obj;
    size_t pubkey_len;

    uint64_t local_attempts = 0;

    // Inizializza la chiave privata con un valore casuale nel range del thread
    init_random_privkey_in_range(privkey, &seed, data->range_start, data->range_end);

    // Mostra la chiave privata di partenza per questo thread
    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 ultra-ottimizzato con prefetching e branch reduction
    pubkey_len = 65; // Costante, settato una volta sola

    while (keep_running) {
        // Processa batch di chiavi consecutive
        for (int batch = 0; batch < BATCH_SIZE; batch++) {
            // Genera chiave pubblica non compressa usando secp256k1
            // Questa è l'operazione più costosa (~95% del tempo)
            if (__builtin_expect(secp256k1_ec_pubkey_create(ctx, &pubkey_obj, privkey), 1)) {
                // Serializza in formato non compresso (65 bytes)
                secp256k1_ec_pubkey_serialize(ctx, pubkey, &pubkey_len,
                                             &pubkey_obj, SECP256K1_EC_UNCOMPRESSED);

                // Verifica corrispondenza (Bloom filter first = velocissimo)
                // Solo ~0.001% dei casi passerà il Bloom filter
                if (__builtin_expect(check_match(pubkey), 0)) {
                    save_found_key(privkey, pubkey);
                }
            }

            // Incrementa la chiave privata di 1 (inline, operazioni a 64-bit)
            increment_privkey(privkey);
        }

        local_attempts += BATCH_SIZE;

        // Aggiorna contatore globale (senza lock - ogni thread scrive solo il proprio indice)
        attempts_per_thread[thread_id] = local_attempts;

        // Check keep_running solo una volta per batch invece che ad ogni iterazione
        if (__builtin_expect(!keep_running, 0)) break;
    }

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

int main(int argc, char** argv) {
    printf("========================================\n");
    printf("  Bitcoin P2PK Bruteforce v1.0\n");
    printf("  SOLO PER SCOPI EDUCATIVI\n");
    printf("========================================\n\n");

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

    // Inizializza secp256k1 con flag ottimizzato per verifiche multiple
    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 il contesto per migliorare la sicurezza e performance
    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] Impossibile randomizzare contesto secp256k1\n");
            }
        } else {
            fprintf(stderr, "[WARNING] Impossibile leggere entropy da /dev/urandom\n");
        }
    }

    // Carica chiavi target
    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 disponibili
    num_threads = get_num_threads();
    printf("[+] CPU rilevata: %d thread disponibili\n", num_threads);
    printf("[+] Partizionamento spazio chiavi in %d regioni\n", num_threads);

    // Inizializza timestamp e seed base robusto
    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;

        // Seed molto distanziati: combina timestamp, thread_id e random
        // Questo garantisce seed completamente diversi anche se lanciato rapidamente
        uint64_t base_seed = (uint64_t)time(NULL);
        uint64_t thread_offset = ((uint64_t)i << 48);  // Usa i bit alti
        uint64_t random_part = ((uint64_t)rand() << 32) | rand();
        thread_data[i].seed = base_seed ^ thread_offset ^ random_part;

        // Partiziona lo spazio delle chiavi
        partition_keyspace(i, num_threads, thread_data[i].range_start, thread_data[i].range_end);

        // Mostra info del range (primi 4 bytes per brevità)
        printf("  Thread %d: range 0x%02x%02x%02x%02x... - 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].range_end[0], thread_data[i].range_end[1],
               thread_data[i].range_end[2], thread_data[i].range_end[3],
               thread_data[i].seed);

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

    printf("\n");

    // Loop principale - mostra progresso
    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;
}