//! kHeavyHash Proof of Work algorithm.
//!
//! The kHeavyHash algorithm is designed to be ASIC-resistant by combining:
//! - SHA3-256 for cryptographic hashing
//! - Matrix multiplication in GF(2^8) for memory hardness
//! - Blake3 for fast nonce mixing
//!
//! Algorithm steps:
//! 1. pre_hash = SHA3-256(header_without_nonce)
//! 2. nonce_hash = Blake3(pre_hash || nonce)
//! 3. matrix_out = HeavyMatrix × pre_hash (in GF(2^8))
//! 4. mixed = matrix_out XOR nonce_hash
//! 5. pow_hash = SHA3-256(mixed)

use sha3::{Digest, Sha3_256};
use synor_types::Hash256;

use crate::matrix::OptimizedMatrix;
use crate::Target;

/// The kHeavyHash algorithm implementation.
pub struct KHeavyHash {
    /// The optimized matrix for fast multiplication.
    matrix: OptimizedMatrix,
}

impl KHeavyHash {
    /// Creates a new kHeavyHash instance.
    pub fn new() -> Self {
        KHeavyHash {
            matrix: OptimizedMatrix::new(),
        }
    }

    /// Creates with a custom matrix seed.
    pub fn with_seed(seed: &[u8; 32]) -> Self {
        KHeavyHash {
            matrix: OptimizedMatrix::from_seed(seed),
        }
    }

    /// Computes the full kHeavyHash PoW.
    ///
    /// # Arguments
    /// * `header` - Block header bytes (without nonce)
    /// * `nonce` - The nonce to try
    ///
    /// # Returns
    /// The PoW hash that must be compared against target.
    pub fn hash(&self, header: &[u8], nonce: u64) -> PowHash {
        // Step 1: Pre-hash the header
        let pre_hash = self.pre_hash(header);

        // Step 2-5: Compute final PoW hash with nonce
        self.finalize(&pre_hash, nonce)
    }

    /// Pre-hashes the block header (can be cached).
    ///
    /// This is SHA3-256 of the header without the nonce.
    pub fn pre_hash(&self, header: &[u8]) -> Hash256 {
        let mut hasher = Sha3_256::new();
        hasher.update(header);
        let result: [u8; 32] = hasher.finalize().into();
        Hash256::from_bytes(result)
    }

    /// Finalizes the PoW hash given a pre-hash and nonce.
    ///
    /// This allows mining to cache the pre-hash and only vary the nonce.
    pub fn finalize(&self, pre_hash: &Hash256, nonce: u64) -> PowHash {
        // Step 2: Hash nonce with pre_hash using Blake3
        let nonce_bytes = nonce.to_le_bytes();
        let mut nonce_input = [0u8; 40]; // 32 + 8
        nonce_input[..32].copy_from_slice(pre_hash.as_bytes());
        nonce_input[32..].copy_from_slice(&nonce_bytes);
        let nonce_hash: [u8; 32] = blake3::hash(&nonce_input).into();

        // Step 3: Matrix multiplication
        let matrix_out = self.matrix.multiply(pre_hash.as_bytes());

        // Step 4: XOR matrix output with nonce hash
        let mut mixed = [0u8; 32];
        for i in 0..32 {
            mixed[i] = matrix_out[i] ^ nonce_hash[i];
        }

        // Step 5: Final SHA3-256
        let mut hasher = Sha3_256::new();
        hasher.update(mixed);
        let final_hash: [u8; 32] = hasher.finalize().into();

        PowHash {
            hash: Hash256::from_bytes(final_hash),
            nonce,
            pre_hash: *pre_hash,
        }
    }

    /// Verifies a PoW solution.
    pub fn verify(&self, header: &[u8], nonce: u64, target: &Target) -> bool {
        let pow = self.hash(header, nonce);
        target.is_met_by(&pow.hash)
    }

    /// Mines a block by trying nonces until target is met.
    ///
    /// # Arguments
    /// * `header` - Block header bytes
    /// * `target` - Difficulty target
    /// * `start_nonce` - Starting nonce
    /// * `max_tries` - Maximum nonces to try (0 = unlimited)
    ///
    /// # Returns
    /// The solving nonce and PoW hash, or None if not found.
    pub fn mine(
        &self,
        header: &[u8],
        target: &Target,
        start_nonce: u64,
        max_tries: u64,
    ) -> Option<PowHash> {
        let pre_hash = self.pre_hash(header);
        let limit = if max_tries == 0 {
            u64::MAX
        } else {
            start_nonce.saturating_add(max_tries)
        };

        for nonce in start_nonce..limit {
            let pow = self.finalize(&pre_hash, nonce);
            if target.is_met_by(&pow.hash) {
                return Some(pow);
            }
        }

        None
    }

    /// Mines with a callback for progress reporting.
    pub fn mine_with_callback<F>(
        &self,
        header: &[u8],
        target: &Target,
        start_nonce: u64,
        max_tries: u64,
        mut callback: F,
    ) -> Option<PowHash>
    where
        F: FnMut(u64, u64) -> bool, // (nonces_tried, current_nonce) -> continue?
    {
        let pre_hash = self.pre_hash(header);
        let limit = if max_tries == 0 {
            u64::MAX
        } else {
            start_nonce.saturating_add(max_tries)
        };

        let mut tried = 0u64;
        for nonce in start_nonce..limit {
            let pow = self.finalize(&pre_hash, nonce);
            tried += 1;

            if target.is_met_by(&pow.hash) {
                return Some(pow);
            }

            // Report progress every 10000 hashes
            if tried % 10000 == 0 && !callback(tried, nonce) {
                return None; // Cancelled
            }
        }

        None
    }

    /// Returns the matrix used by this hasher.
    pub fn matrix(&self) -> &OptimizedMatrix {
        &self.matrix
    }
}

impl Default for KHeavyHash {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for KHeavyHash {
    fn clone(&self) -> Self {
        // Matrix is large, but we need to clone for multi-threading
        KHeavyHash {
            matrix: OptimizedMatrix::new(), // Recreate from default seed
        }
    }
}

/// Result of a PoW computation.
#[derive(Clone, Debug)]
pub struct PowHash {
    /// The final PoW hash.
    pub hash: Hash256,
    /// The nonce that produced this hash.
    pub nonce: u64,
    /// The pre-hash (for verification).
    pub pre_hash: Hash256,
}

impl PowHash {
    /// Checks if this PoW meets a target.
    pub fn meets_target(&self, target: &Target) -> bool {
        target.is_met_by(&self.hash)
    }

    /// Returns the hash as bytes.
    pub fn as_bytes(&self) -> &[u8; 32] {
        self.hash.as_bytes()
    }
}

/// Parallel mining using multiple threads.
pub struct ParallelMiner {
    /// kHeavyHash instances per thread.
    hashers: Vec<KHeavyHash>,
    /// Number of threads.
    num_threads: usize,
}

impl ParallelMiner {
    /// Creates a parallel miner with the specified number of threads.
    pub fn new(num_threads: usize) -> Self {
        let threads = if num_threads == 0 {
            num_cpus()
        } else {
            num_threads
        };

        let hashers = (0..threads).map(|_| KHeavyHash::new()).collect();

        ParallelMiner {
            hashers,
            num_threads: threads,
        }
    }

    /// Returns the number of threads.
    pub fn num_threads(&self) -> usize {
        self.num_threads
    }

    /// Mines using all threads.
    ///
    /// Each thread gets a different nonce range to search.
    pub fn mine(&self, header: &[u8], target: &Target, start_nonce: u64) -> Option<PowHash> {
        use std::sync::atomic::{AtomicBool, Ordering};
        use std::sync::Arc;

        let found = Arc::new(AtomicBool::new(false));
        let result = Arc::new(parking_lot::Mutex::new(None));

        let nonce_step = u64::MAX / self.num_threads as u64;

        std::thread::scope(|s| {
            for (i, hasher) in self.hashers.iter().enumerate() {
                let thread_start = start_nonce.wrapping_add(i as u64 * nonce_step);
                let found = Arc::clone(&found);
                let result = Arc::clone(&result);
                let header = header.to_vec();
                let target = *target;

                s.spawn(move || {
                    let pre_hash = hasher.pre_hash(&header);
                    let mut nonce = thread_start;

                    loop {
                        // Check if another thread found it
                        if found.load(Ordering::Relaxed) {
                            return;
                        }

                        let pow = hasher.finalize(&pre_hash, nonce);

                        if target.is_met_by(&pow.hash) {
                            found.store(true, Ordering::Relaxed);
                            *result.lock() = Some(pow);
                            return;
                        }

                        nonce = nonce.wrapping_add(1);

                        // Don't wrap into another thread's range
                        if nonce == thread_start.wrapping_add(nonce_step) {
                            return;
                        }
                    }
                });
            }
        });

        Arc::try_unwrap(result).ok()?.into_inner()
    }
}

/// Returns the number of CPU cores.
fn num_cpus() -> usize {
    std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1)
}

/// Benchmark helper for hashrate measurement.
pub struct HashrateBenchmark {
    hasher: KHeavyHash,
}

impl HashrateBenchmark {
    /// Creates a new benchmark.
    pub fn new() -> Self {
        HashrateBenchmark {
            hasher: KHeavyHash::new(),
        }
    }

    /// Runs a benchmark for the specified duration.
    ///
    /// Returns hashes per second.
    pub fn run(&self, duration_ms: u64) -> f64 {
        let header = [0x42u8; 80]; // Dummy header
        let pre_hash = self.hasher.pre_hash(&header);

        let start = std::time::Instant::now();
        let mut count = 0u64;
        let mut nonce = 0u64;

        while start.elapsed().as_millis() < duration_ms as u128 {
            let _ = self.hasher.finalize(&pre_hash, nonce);
            nonce = nonce.wrapping_add(1);
            count += 1;
        }

        let elapsed = start.elapsed().as_secs_f64();
        count as f64 / elapsed
    }
}

impl Default for HashrateBenchmark {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_kheavyhash_deterministic() {
        let hasher = KHeavyHash::new();
        let header = b"test block header data";
        let nonce = 12345u64;

        let hash1 = hasher.hash(header, nonce);
        let hash2 = hasher.hash(header, nonce);

        assert_eq!(hash1.hash, hash2.hash);
        assert_eq!(hash1.nonce, hash2.nonce);
    }

    #[test]
    fn test_kheavyhash_different_nonces() {
        let hasher = KHeavyHash::new();
        let header = b"test block header data";

        let hash1 = hasher.hash(header, 1);
        let hash2 = hasher.hash(header, 2);

        assert_ne!(hash1.hash, hash2.hash);
    }

    #[test]
    fn test_pre_hash_caching() {
        let hasher = KHeavyHash::new();
        let header = b"test block header data";

        let pre_hash = hasher.pre_hash(header);

        let pow1 = hasher.finalize(&pre_hash, 1);
        let pow2 = hasher.finalize(&pre_hash, 2);

        assert_eq!(pow1.pre_hash, pow2.pre_hash);
        assert_ne!(pow1.hash, pow2.hash);
    }

    #[test]
    fn test_mine_easy_target() {
        let hasher = KHeavyHash::new();
        let header = b"mine me";

        // Very easy target (almost all 1s)
        // Note: Target::max() still requires first byte to be 0x00,
        // which has ~1/256 probability, so we need enough tries.
        let target = Target::max();

        // 10000 tries gives ~39 expected successes (10000/256)
        let result = hasher.mine(header, &target, 0, 10000);
        assert!(result.is_some());

        let pow = result.unwrap();
        assert!(target.is_met_by(&pow.hash));
    }

    #[test]
    fn test_verify() {
        let hasher = KHeavyHash::new();
        let header = b"verify test";
        let target = Target::max();

        // Need enough tries to find a valid nonce
        let pow = hasher.mine(header, &target, 0, 10000).unwrap();

        assert!(hasher.verify(header, pow.nonce, &target));
        // Wrong nonce should fail (unless extremely lucky)
        // Skip this as it might randomly pass with easy target
    }

    #[test]
    fn test_hashrate_benchmark() {
        let bench = HashrateBenchmark::new();
        let hashrate = bench.run(100); // 100ms benchmark
        assert!(hashrate > 0.0);
        println!("Benchmark: {:.0} H/s", hashrate);
    }

    #[test]
    fn test_parallel_miner() {
        let miner = ParallelMiner::new(2);
        assert!(miner.num_threads() >= 1);

        let header = b"parallel mining test";
        let target = Target::max();

        let result = miner.mine(header, &target, 0);
        assert!(result.is_some());
    }
}
