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

use std::{
    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`.
    RandomXFlags::get_recommended_flags() | RandomXFlags::FULLMEM
}

fn init_dataset(flags: RandomXFlags, input: &[u8; 32]) -> RandomXDataset {
    // Allocate cache and dataset
    let cache = RandomXCache::new(flags, &input[..]).unwrap();
    let dataset_item_count = RandomXDataset::count().unwrap();
    let dataset = RandomXDataset::new(flags, cache, dataset_item_count).unwrap();

    // Multithreaded dataset init using all available threads
    let threads = thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    println!("Initializing RandomX dataset using {threads} threads...");
    let mut handles = Vec::with_capacity(threads);
    let threads_u32 = threads as u32;
    let per_thread = dataset_item_count / threads_u32;
    let remainder = dataset_item_count % threads_u32;
    for t in 0..threads_u32 {
        let dataset = dataset.clone();
        let start_item = t * per_thread;
        let count = per_thread + if t == threads_u32 - 1 { remainder } else { 0 };
        handles.push(thread::spawn(move || {
            dataset.init(start_item, count);
        }));
    }

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

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 dataset = init_dataset(flags, input);
    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!");
    header.nonce = 0;
    let mining_start = Instant::now();
    vm.calculate_hash_first(header.hash().as_bytes()).unwrap();
    loop {
        header.nonce += 1;
        let out_hash = vm.calculate_hash_next(header.hash().as_bytes()).unwrap();
        let out_hash = BigUint::from_bytes_le(&out_hash);
        if &out_hash <= target {
            header.nonce -= 1; // Since out hash refers to previous nonce
            println!("Found block header using nonce {}", header.nonce);
            println!("Block header hash {}", header.hash());
            println!("RandomX output: 0x{out_hash:064x}");
            break;
        }
    }
    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 dataset = init_dataset(flags, input);
    println!("Setup time: {:?}", setup_start.elapsed());

    println!("Initializing mining threads...");
    let mut handles = Vec::with_capacity(threads);
    let atomic_nonce = Arc::new(AtomicU64::new(0));
    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();
        let atomic_nonce = atomic_nonce.clone();
        let found_header = found_header.clone();
        let found_nonce = found_nonce.clone();
        let dataset = dataset.clone();

        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());
            let mut last_nonce = atomic_nonce.fetch_add(1, Ordering::SeqCst);
            thread_header.nonce = last_nonce;
            vm.calculate_hash_first(thread_header.hash().as_bytes())
                .unwrap();
            loop {
                if found_header.load(Ordering::SeqCst) {
                    println!("Block header found, thread #{t} exiting");
                    break;
                }

                thread_header.nonce = atomic_nonce.fetch_add(1, Ordering::SeqCst);
                let out_hash = vm
                    .calculate_hash_next(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);
                    thread_header.nonce = last_nonce; // Since out hash refers to previous run nonce
                    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;
                }
                last_nonce = thread_header.nonce;
            }
        }));
    }

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