//! randomx example that mines a dummy header using multiple threads

use std::{
    collections::VecDeque,
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        Arc,
    },
    thread,
    time::Instant,
};

use num_bigint::BigUint;
use randomx::*;

#[derive(Clone, Debug)]
struct Header {
    version: u8,
    previous: blake3::Hash,
    height: u32,
    nonce: u64,
}

impl Header {
    fn new(previous: blake3::Hash) -> Self {
        Self {
            version: 1,
            previous,
            height: 1,
            nonce: 0,
        }
    }

    fn hash(&self) -> blake3::Hash {
        let mut hasher = blake3::Hasher::new();
        hasher.update(&self.version.to_le_bytes());
        hasher.update(self.previous.as_bytes());
        hasher.update(&self.height.to_le_bytes());
        hasher.update(&self.nonce.to_le_bytes());
        hasher.finalize()
    }
}

fn get_mining_flags() -> RandomXFlags {
    // Try adding `| RandomXFlags::LARGEPAGES`.
    let mut flags = RandomXFlags::get_recommended_flags() | RandomXFlags::FULLMEM;
    if is_x86_feature_detected!("avx2") {
        flags |= RandomXFlags::ARGON2_AVX2;
    } else if is_x86_feature_detected!("ssse3") {
        flags |= RandomXFlags::ARGON2_SSSE3;
    }
    flags
}

fn single_thread_mine(header: &mut Header, target: &BigUint, input: &[u8; 32]) {
    println!("Initializing RandomX cache and dataset...");
    let setup_start = Instant::now();
    let flags = get_mining_flags();
    let cache = RandomXCache::new(flags, &input[..]).unwrap();
    let dataset_item_count = RandomXDataset::count().unwrap();
    let dataset = RandomXDataset::new_init(flags, cache, 0, dataset_item_count).unwrap();
    println!("Setup time: {:?}", setup_start.elapsed());

    println!("Initializing RandomX VM...");
    let vm_start = Instant::now();
    let vm = RandomXVM::new(flags, None, Some(dataset)).unwrap();
    println!("Initialized RandomX VM in {:?}", vm_start.elapsed());

    println!("Mining started!");
    let mining_start = Instant::now();
    loop {
        let out_hash = vm.calculate_hash(header.hash().as_bytes()).unwrap();
        let out_hash = BigUint::from_bytes_le(&out_hash);
        if &out_hash <= target {
            println!("Found block header using nonce {}", header.nonce);
            println!("Block header hash {}", header.hash());
            println!("RandomX output: 0x{out_hash:064x}");
            break;
        }

        header.nonce += 1;
    }
    println!("Completed mining in {:?}", mining_start.elapsed());
    println!("Mined header: {header:?}");
}

fn multi_thread_mine(threads: usize, header: &mut Header, target: &BigUint, input: &[u8; 32]) {
    println!("Initializing RandomX cache and dataset...");
    let setup_start = Instant::now();
    let flags = get_mining_flags();
    let cache = RandomXCache::new(flags, &input[..]).unwrap();
    let dataset_item_count = RandomXDataset::count().unwrap();
    let dataset = RandomXDataset::new(flags, cache, dataset_item_count).unwrap();
    let mut subsets = VecDeque::with_capacity(threads);

    // Multithreaded dataset init
    let threads_u32 = threads as u32;
    for t in 0..threads_u32 {
        println!("Initializing RandomX dataset for thread #{t}...");
        let ds_start = Instant::now();
        let a = (dataset_item_count * t) / threads_u32;
        let b = (dataset_item_count * (t + 1)) / threads_u32;
        subsets.push_back(dataset.subset_init(a, b - a));
        println!(
            "Initialized RandomX dataset for thread #{t} in {:?}",
            ds_start.elapsed()
        );
    }
    println!("Setup time: {:?}", setup_start.elapsed());

    println!("Initializing mining threads...");
    let mut handles = Vec::with_capacity(threads);
    let found_header = Arc::new(AtomicBool::new(false));
    let found_nonce = Arc::new(AtomicU64::new(0));
    let threads_u64 = threads as u64;
    let mining_start = Instant::now();
    for t in 0..threads_u64 {
        if found_header.load(Ordering::SeqCst) {
            println!("Block header found, threads creation loop exiting");
            break;
        }

        let target = target.clone();
        let mut thread_header = header.clone();
        thread_header.nonce = t;
        let found_header = Arc::clone(&found_header);
        let found_nonce = Arc::clone(&found_nonce);
        let dataset = subsets.pop_front().unwrap();

        handles.push(thread::spawn(move || {
            println!("Initializing RandomX VM #{t}...");
            let vm_start = Instant::now();
            let vm = RandomXVM::new(flags, None, Some(dataset)).unwrap();
            println!("Initialized RandomX VM #{t} in {:?}", vm_start.elapsed());
            loop {
                if found_header.load(Ordering::SeqCst) {
                    println!("Block header found, thread #{t} exiting");
                    break;
                }

                let out_hash = vm.calculate_hash(thread_header.hash().as_bytes()).unwrap();
                let out_hash = BigUint::from_bytes_le(&out_hash);
                if out_hash <= target {
                    found_header.store(true, Ordering::SeqCst);
                    found_nonce.store(thread_header.nonce, Ordering::SeqCst);
                    println!(
                        "Thread #{t} found block header using nonce {}",
                        thread_header.nonce
                    );
                    println!("Block header hash {}", thread_header.hash());
                    println!("RandomX output: 0x{out_hash:064x}");
                    break;
                }

                // This means thread 0 will use nonces, 0, 4, 8, ...
                // and thread 1 will use nonces, 1, 5, 9, ...
                thread_header.nonce += threads_u64;
            }
        }));
    }

    // Wait for threads to finish mining
    for handle in handles {
        let _ = handle.join();
    }

    println!("Completed mining in {:?}", mining_start.elapsed());
    header.nonce = found_nonce.load(Ordering::SeqCst);
    println!("Mined header: {header:?}");
}

fn main() {
    const THREADS: usize = 8;
    let difficulty = BigUint::from(2500_u32);
    let target = BigUint::from_bytes_le(&[0xFF; 32]) / difficulty;
    let previous = blake3::hash(b"Let there be dark!");
    let input = previous.as_bytes();
    let mut header = Header::new(previous);
    println!("Mine target: 0x{target:064x}");
    println!("PoW input: {previous}");

    match THREADS {
        0 => panic!("Can't use 0 threads!"),
        1 => single_thread_mine(&mut header, &target, input),
        _ => multi_thread_mine(THREADS, &mut header, &target, input),
    }

    println!("Initializing verifier...");
    let verifier_setup = Instant::now();
    let flags = RandomXFlags::get_recommended_flags();
    let cache = RandomXCache::new(flags, input).unwrap();
    let vm = RandomXVM::new(flags, Some(cache), None).unwrap();
    println!("Setup time: {:?}", verifier_setup.elapsed());

    let hashing_time = Instant::now();
    let out_hash = vm.calculate_hash(header.hash().as_bytes()).unwrap();
    println!("Completed hashing in {:?}", hashing_time.elapsed());
    let out_hash = BigUint::from_bytes_le(&out_hash);
    println!("RandomX output: 0x{out_hash:064x}");
    assert!(out_hash <= target);
}