use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Instant;

use randomx::*;

const BLOCK_TEMPLATE: [u8; 76] = [
    0x07, 0x07, 0xf7, 0xa4, 0xf0, 0xd6, 0x05, 0xb3, 0x03, 0x26, 0x08, 0x16, 0xba, 0x3f, 0x10, 0x90,
    0x2e, 0x1a, 0x14, 0x5a, 0xc5, 0xfa, 0xd3, 0xaa, 0x3a, 0xf6, 0xea, 0x44, 0xc1, 0x18, 0x69, 0xdc,
    0x4f, 0x85, 0x3f, 0x00, 0x2b, 0x2e, 0xea, 0x00, 0x00, 0x00, 0x00, 0x77, 0xb2, 0x06, 0xa0, 0x2c,
    0xa5, 0xb1, 0xd4, 0xce, 0x6b, 0xbf, 0xdf, 0x0a, 0xca, 0xc3, 0x8b, 0xde, 0xd3, 0x4d, 0x2d, 0xcd,
    0xee, 0xf9, 0x5c, 0xd2, 0x0c, 0xef, 0xc1, 0x2f, 0x61, 0xd5, 0x61, 0x09,
];

struct AtomicHash {
    hash: [AtomicU64; 4],
}

impl AtomicHash {
    fn new() -> Self {
        Self {
            hash: [
                AtomicU64::new(0),
                AtomicU64::new(0),
                AtomicU64::new(0),
                AtomicU64::new(0),
            ],
        }
    }

    fn xor_with(&self, update: &[u64; 4]) {
        for (i, hash) in self.hash.iter().enumerate() {
            hash.fetch_xor(update[i], Ordering::SeqCst);
        }
    }

    fn print(&self) {
        for i in 0..4 {
            let h = self.hash[i].load(Ordering::SeqCst);
            let bytes = h.to_le_bytes();
            for byte in bytes {
                print!("{:02x}", byte);
            }
        }
        println!();
    }
}

fn mine<const BATCH: bool, const COMMIT: bool>(
    vm: &RandomXVM,
    atomic_nonce: &AtomicU32,
    result: &AtomicHash,
    nonces_count: u32,
    _thread: usize,
    _cpuid: i32,
) {
    let mut block_template = BLOCK_TEMPLATE;
    let mut nonce = atomic_nonce.fetch_add(1, Ordering::SeqCst);

    if BATCH {
        block_template[39..43].copy_from_slice(&nonce.to_le_bytes());
        vm.calculate_hash_first(&block_template).unwrap();
    }

    while nonce < nonces_count {
        if BATCH {
            nonce = atomic_nonce.fetch_add(1, Ordering::SeqCst);
        }

        block_template[39..43].copy_from_slice(&nonce.to_le_bytes());

        let mut hash = if BATCH {
            vm.calculate_hash_next(&block_template).unwrap()
        } else {
            vm.calculate_hash(&block_template).unwrap()
        };

        if COMMIT {
            let mut commitment = vec![0u8; RANDOMX_HASH_SIZE as usize];
            calculate_commitment(&block_template, &hash, &mut commitment).unwrap();
            hash = commitment;
        }

        let mut hash_u64 = [0u64; 4];
        for i in 0..4 {
            hash_u64[i] = u64::from_le_bytes([
                hash[i * 8],
                hash[i * 8 + 1],
                hash[i * 8 + 2],
                hash[i * 8 + 3],
                hash[i * 8 + 4],
                hash[i * 8 + 5],
                hash[i * 8 + 6],
                hash[i * 8 + 7],
            ]);
        }

        result.xor_with(&hash_u64);

        if !BATCH {
            nonce = atomic_nonce.fetch_add(1, Ordering::SeqCst);
        }
    }
}

fn print_usage(executable: &str) {
    println!("Usage: {} [OPTIONS]", executable);
    println!("Supported options:");
    println!("  --help        shows this message");
    println!("  --mine        mining mode: 2080 MiB");
    println!("  --verify      verification mode: 256 MiB");
    println!("  --jit         JIT compiled mode (default: interpreter)");
    println!("  --secure      W^X policy for JIT pages (default: off)");
    println!("  --largePages  use large pages (default: small pages)");
    println!("  --softAes     use software AES (default: hardware AES)");
    println!("  --threads T   use T threads (default: 1)");
    println!("  --affinity A  thread affinity bitmask (default: 0)");
    println!("  --init Q      initialize dataset with Q threads (default: 1)");
    println!("  --nonces N    run N nonces (default: 1000)");
    println!("  --seed S      seed for cache initialization (default: 0)");
    println!("  --ssse3       use optimized Argon2 for SSSE3 CPUs");
    println!("  --avx2        use optimized Argon2 for AVX2 CPUs");
    println!("  --auto        select the best options for the current CPU");
    println!("  --noBatch     calculate hashes one by one (default: batch)");
    println!("  --commit      calculate commitments instead of hashes (default: hashes)");
}

fn has_arg(args: &[String], name: &str) -> bool {
    args.iter().any(|arg| arg == name)
}

fn get_int_arg(args: &[String], name: &str, default: i32) -> i32 {
    args.iter()
        .position(|arg| arg == name)
        .and_then(|pos| args.get(pos + 1))
        .and_then(|val| val.parse().ok())
        .unwrap_or(default)
}

fn get_u64_arg(args: &[String], name: &str, default: u64) -> u64 {
    args.iter()
        .position(|arg| arg == name)
        .and_then(|pos| args.get(pos + 1))
        .and_then(|val| val.parse().ok())
        .unwrap_or(default)
}

fn cpuid_from_mask(mask: u64, index: usize) -> i32 {
    let mut count = 0;
    for i in 0..64 {
        if (mask & (1u64 << i)) != 0 {
            if count == index {
                return i;
            }
            count += 1;
        }
    }
    -1
}

fn mask_to_string(mask: u64) -> String {
    let mut result = String::new();
    let mut first = true;
    for i in 0..64 {
        if (mask & (1u64 << i)) != 0 {
            if !first {
                result.push(',');
            }
            result.push_str(&i.to_string());
            first = false;
        }
    }
    result
}

fn main() {
    let args: Vec<String> = std::env::args().collect();

    let help = has_arg(&args, "--help");
    let mining_mode = has_arg(&args, "--mine");
    let verification_mode = has_arg(&args, "--verify");
    let soft_aes = has_arg(&args, "--softAes");
    let large_pages = has_arg(&args, "--largePages") || has_arg(&args, "--largepages");
    let jit = has_arg(&args, "--jit");
    let secure = has_arg(&args, "--secure");
    let ssse3 = has_arg(&args, "--ssse3");
    let avx2 = has_arg(&args, "--avx2");
    let auto_flags = has_arg(&args, "--auto");
    let no_batch = has_arg(&args, "--noBatch");
    let commit = has_arg(&args, "--commit");

    let thread_count = get_int_arg(&args, "--threads", 1) as usize;
    let thread_affinity = get_u64_arg(&args, "--affinity", 0);
    let nonces_count = get_int_arg(&args, "--nonces", 1000) as u32;
    let mut init_thread_count = get_int_arg(&args, "--init", 1) as usize;
    let seed_value = get_int_arg(&args, "--seed", 0);

    let seed = seed_value.to_le_bytes();

    println!("RandomX benchmark v1.2.1");

    if help {
        print_usage(&args[0]);
        return;
    }

    if !mining_mode && !verification_mode {
        println!("Please select either the fast mode (--mine) or the slow mode (--verify)");
        println!("Run '{}' --help' to see all supported options", args[0]);
        return;
    }

    let mut flags = if auto_flags {
        init_thread_count = thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);
        RandomXFlags::get_recommended_flags()
    } else {
        let mut flags = RandomXFlags::DEFAULT;
        if ssse3 {
            flags |= RandomXFlags::ARGON2_SSSE3;
        }
        if avx2 {
            flags |= RandomXFlags::ARGON2_AVX2;
        }
        if !soft_aes {
            flags |= RandomXFlags::HARDAES;
        }
        if jit {
            flags |= RandomXFlags::JIT;
        }
        flags
    };

    if large_pages {
        flags |= RandomXFlags::LARGEPAGES;
    }
    if mining_mode {
        flags |= RandomXFlags::FULLMEM;
    }
    if secure {
        flags |= RandomXFlags::SECURE;
    }

    // Print configuration
    if flags.contains(RandomXFlags::ARGON2_AVX2) {
        println!(" - Argon2 implementation: AVX2");
    } else if flags.contains(RandomXFlags::ARGON2_SSSE3) {
        println!(" - Argon2 implementation: SSSE3");
    } else {
        println!(" - Argon2 implementation: reference");
    }

    if flags.contains(RandomXFlags::FULLMEM) {
        println!(" - full memory mode (2080 MiB)");
    } else {
        println!(" - light memory mode (256 MiB)");
    }

    if flags.contains(RandomXFlags::JIT) {
        print!(" - JIT compiled mode");
        if flags.contains(RandomXFlags::SECURE) {
            print!(" (secure)");
        }
        println!();
    } else {
        println!(" - interpreted mode");
    }

    if flags.contains(RandomXFlags::HARDAES) {
        println!(" - hardware AES mode");
    } else {
        println!(" - software AES mode");
    }

    if flags.contains(RandomXFlags::LARGEPAGES) {
        println!(" - large pages mode");
    } else {
        println!(" - small pages mode");
    }

    if thread_affinity != 0 {
        println!(" - thread affinity ({})", mask_to_string(thread_affinity));
    }

    if no_batch {
        if commit {
            println!(" - hash commitments");
        }
    } else if commit {
        println!(" - hash commitments");
    } else {
        println!(" - batch mode");
    }

    print!("Initializing");
    if mining_mode {
        print!(
            " ({} thread{})",
            init_thread_count,
            if init_thread_count > 1 { "s" } else { "" }
        );
    }
    println!(" ...");

    let start = Instant::now();

    let cache = RandomXCache::new(flags, &seed).unwrap();

    let dataset = if mining_mode {
        let dataset_item_count = RandomXDataset::count().unwrap();
        let dataset = RandomXDataset::new(flags, cache.clone(), dataset_item_count).unwrap();

        if init_thread_count > 1 {
            let per_thread = dataset_item_count / init_thread_count as u32;
            let remainder = dataset_item_count % init_thread_count as u32;
            let mut handles = vec![];

            for i in 0..init_thread_count {
                let dataset_clone = dataset.clone();
                let start_item = i as u32 * per_thread;
                let count = per_thread
                    + if i == init_thread_count - 1 {
                        remainder
                    } else {
                        0
                    };

                let handle = thread::spawn(move || {
                    dataset_clone.init(start_item, count);
                });
                handles.push(handle);
            }

            for handle in handles {
                handle.join().unwrap();
            }
        } else {
            dataset.init(0, dataset_item_count);
        }

        Some(dataset)
    } else {
        None
    };

    let elapsed = start.elapsed();
    println!("Memory initialized in {:.3} s", elapsed.as_secs_f64());

    println!("Initializing {} virtual machine(s) ...", thread_count);
    let mut vms = Vec::new();
    for _ in 0..thread_count {
        let vm = if mining_mode {
            RandomXVM::new(flags, None, dataset.clone()).unwrap()
        } else {
            RandomXVM::new(flags, Some(cache.clone()), None).unwrap()
        };
        vms.push(Arc::new(vm));
    }

    println!("Running benchmark ({} nonces) ...", nonces_count);
    let atomic_nonce = Arc::new(AtomicU32::new(0));
    let result = Arc::new(AtomicHash::new());

    let start = Instant::now();

    if thread_count > 1 {
        let mut handles = vec![];
        for (i, vm) in vms.iter().enumerate() {
            let vm = vm.clone();
            let atomic_nonce = atomic_nonce.clone();
            let result = result.clone();
            let cpuid = if thread_affinity != 0 {
                cpuid_from_mask(thread_affinity, i)
            } else {
                -1
            };

            let handle = thread::spawn(move || match (no_batch, commit) {
                (false, false) => {
                    mine::<true, false>(&vm, &atomic_nonce, &result, nonces_count, i, cpuid)
                }
                (false, true) => {
                    mine::<false, true>(&vm, &atomic_nonce, &result, nonces_count, i, cpuid)
                }
                (true, false) => {
                    mine::<false, false>(&vm, &atomic_nonce, &result, nonces_count, i, cpuid)
                }
                (true, true) => {
                    mine::<false, true>(&vm, &atomic_nonce, &result, nonces_count, i, cpuid)
                }
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.join().unwrap();
        }
    } else {
        match (no_batch, commit) {
            (false, false) => {
                mine::<true, false>(&vms[0], &atomic_nonce, &result, nonces_count, 0, -1)
            }
            (false, true) => {
                mine::<false, true>(&vms[0], &atomic_nonce, &result, nonces_count, 0, -1)
            }
            (true, false) => {
                mine::<false, false>(&vms[0], &atomic_nonce, &result, nonces_count, 0, -1)
            }
            (true, true) => {
                mine::<false, true>(&vms[0], &atomic_nonce, &result, nonces_count, 0, -1)
            }
        }
    }

    let elapsed = start.elapsed();

    print!("Calculated result: ");
    result.print();
    if nonces_count == 1000 && seed_value == 0 && !commit {
        println!(
            "Reference result:  10b649a3f15c7c7f88277812f2e74b337a0f20ce909af09199cccb960771cfa1"
        );
    }

    if !mining_mode {
        println!(
            "Performance: {:.3} ms per hash",
            elapsed.as_secs_f64() * 1000.0 / nonces_count as f64
        );
    } else {
        println!(
            "Performance: {:.2} hashes per second",
            nonces_count as f64 / elapsed.as_secs_f64()
        );
    }
}