//! Proof-of-Work (PoW) for admission control.
//!
//! New agents must solve BLAKE3-based puzzles before their assertions are accepted.
//! The difficulty is graduated based on assertion count:
//!
//! - First 10 assertions: 16 bits (~65K iterations, ~16 seconds)
//! - Assertions 11-50: 1 bit (trivial)
//! - 50+ assertions OR trust > 0.6: 0 bits (exempt)
//!
//! # Puzzle Format
//!
//! The agent must find a nonce such that:
//! `BLAKE3(nonce || agent_id || timestamp)` has `difficulty` leading zero bits.
//!
//! # Security Properties
//!
//! - Timestamp prevents replay attacks (max age: 5 minutes)
//! - Agent ID binds proof to specific identity
//! - BLAKE3 provides cryptographic security
//! - Asymmetric cost: O(2^difficulty) to solve, O(1) to verify

use thiserror::Error;

/// Maximum age of a PoW proof in seconds (5 minutes).
pub const POW_MAX_AGE_SECONDS: u64 = 300;

/// Default difficulty for first 10 assertions (16 bits).
pub const POW_INITIAL_DIFFICULTY: u8 = 16;

/// Reduced difficulty for assertions 11-50 (1 bit).
pub const POW_REDUCED_DIFFICULTY: u8 = 1;

/// Threshold for initial difficulty (first 10 assertions).
pub const POW_INITIAL_THRESHOLD: u64 = 10;

/// Threshold for graduation (50 assertions = exempt).
pub const POW_GRADUATED_THRESHOLD: u64 = 50;

/// Errors that can occur during PoW verification.
#[derive(Debug, Clone, Error, PartialEq, Eq)]
pub enum PowError {
    /// The proof timestamp is too old.
    #[error("Proof timestamp expired (max age: {max_age}s, actual age: {actual_age}s)")]
    TimestampExpired {
        /// Maximum allowed age in seconds.
        max_age: u64,
        /// Actual age of the proof in seconds.
        actual_age: u64,
    },

    /// The proof timestamp is in the future.
    #[error(
        "Proof timestamp is in the future (server time: {server_time}, proof time: {proof_time})"
    )]
    TimestampInFuture {
        /// Server's current time.
        server_time: u64,
        /// Proof's timestamp.
        proof_time: u64,
    },

    /// The proof does not meet the required difficulty.
    #[error("Insufficient difficulty (required: {required} leading zeros, found: {found})")]
    InsufficientDifficulty {
        /// Required number of leading zero bits.
        required: u8,
        /// Actual number of leading zero bits.
        found: u8,
    },

    /// The agent ID in the proof does not match the request.
    #[error("Agent ID mismatch")]
    AgentIdMismatch,
}

/// A proof-of-work solution for admission control.
///
/// The proof demonstrates computational effort by finding a nonce that
/// produces a BLAKE3 hash with the required number of leading zero bits.
///
/// # Wire Format
///
/// When transmitted over HTTP, the proof is split into headers:
/// - `X-PoW-Nonce`: The nonce as a decimal string
/// - `X-PoW-Timestamp`: Unix timestamp as a decimal string
/// - `X-Agent-Id`: Agent's Ed25519 public key (hex, existing header)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PowProof {
    /// The nonce value found by the agent.
    pub nonce: u64,
    /// The agent's Ed25519 public key.
    pub agent_id: [u8; 32],
    /// Unix timestamp when the proof was generated.
    pub timestamp: u64,
}

impl PowProof {
    /// Create a new PoW proof.
    #[must_use]
    pub fn new(nonce: u64, agent_id: [u8; 32], timestamp: u64) -> Self {
        Self { nonce, agent_id, timestamp }
    }

    /// Compute the BLAKE3 hash of this proof.
    ///
    /// Hash input: `nonce (8 bytes LE) || agent_id (32 bytes) || timestamp (8 bytes LE)`
    #[must_use]
    pub fn compute_hash(&self) -> blake3::Hash {
        let mut hasher = blake3::Hasher::new();
        hasher.update(&self.nonce.to_le_bytes());
        hasher.update(&self.agent_id);
        hasher.update(&self.timestamp.to_le_bytes());
        hasher.finalize()
    }

    /// Count the number of leading zero bits in a hash.
    ///
    /// # Example
    /// ```
    /// use stemedb_core::types::PowProof;
    ///
    /// let hash = blake3::hash(b"test");
    /// let zeros = PowProof::leading_zeros(&hash);
    /// // Will vary based on hash value
    /// ```
    #[must_use]
    pub fn leading_zeros(hash: &blake3::Hash) -> u8 {
        let bytes = hash.as_bytes();
        let mut count: u8 = 0;

        for byte in bytes {
            if *byte == 0 {
                count = count.saturating_add(8);
            } else {
                count = count.saturating_add(byte.leading_zeros() as u8);
                break;
            }
        }

        count
    }

    /// Verify this proof against the required difficulty.
    ///
    /// # Arguments
    /// * `difficulty` - Required number of leading zero bits
    /// * `max_age` - Maximum allowed age in seconds
    /// * `server_time` - Server's current Unix timestamp
    ///
    /// # Returns
    /// `Ok(())` if the proof is valid, or an error describing the failure.
    ///
    /// # Example
    /// ```
    /// use stemedb_core::types::{PowProof, PowError, POW_MAX_AGE_SECONDS};
    ///
    /// let agent_id = [0u8; 32];
    /// let now = 1700000000u64;
    /// let proof = PowProof::new(12345, agent_id, now);
    ///
    /// // Verification with difficulty 0 always passes (if timestamp is valid)
    /// let result = proof.verify(0, POW_MAX_AGE_SECONDS, now);
    /// assert!(result.is_ok());
    /// ```
    pub fn verify(&self, difficulty: u8, max_age: u64, server_time: u64) -> Result<(), PowError> {
        // Check timestamp is not in the future (with small tolerance for clock skew)
        const CLOCK_SKEW_TOLERANCE: u64 = 30; // 30 seconds
        if self.timestamp > server_time.saturating_add(CLOCK_SKEW_TOLERANCE) {
            return Err(PowError::TimestampInFuture { server_time, proof_time: self.timestamp });
        }

        // Check timestamp is not too old
        let age = server_time.saturating_sub(self.timestamp);
        if age > max_age {
            return Err(PowError::TimestampExpired { max_age, actual_age: age });
        }

        // For difficulty 0, no hash check needed (exempt)
        if difficulty == 0 {
            return Ok(());
        }

        // Compute hash and check leading zeros
        let hash = self.compute_hash();
        let zeros = Self::leading_zeros(&hash);

        if zeros < difficulty {
            return Err(PowError::InsufficientDifficulty { required: difficulty, found: zeros });
        }

        Ok(())
    }

    /// Solve a PoW puzzle by brute-force search.
    ///
    /// This is a convenience method for clients to find a valid nonce.
    /// It iterates from 0 until finding a nonce that satisfies the difficulty.
    ///
    /// # Arguments
    /// * `agent_id` - The agent's Ed25519 public key
    /// * `timestamp` - Unix timestamp for the proof
    /// * `difficulty` - Required number of leading zero bits
    ///
    /// # Returns
    /// A valid `PowProof` with the found nonce.
    ///
    /// # Panics
    /// In theory could run forever if difficulty is impossibly high,
    /// but in practice 16 bits completes in seconds.
    #[must_use]
    pub fn solve(agent_id: [u8; 32], timestamp: u64, difficulty: u8) -> Self {
        // Difficulty 0 means exempt - any nonce works
        if difficulty == 0 {
            return Self::new(0, agent_id, timestamp);
        }

        for nonce in 0..u64::MAX {
            let proof = Self::new(nonce, agent_id, timestamp);
            let hash = proof.compute_hash();
            if Self::leading_zeros(&hash) >= difficulty {
                return proof;
            }
        }

        // Mathematically impossible to reach here with reasonable difficulty
        Self::new(0, agent_id, timestamp)
    }
}

/// Configuration for admission control PoW requirements.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AdmissionConfig {
    /// Difficulty for first `initial_threshold` assertions.
    pub initial_difficulty: u8,
    /// Number of assertions requiring initial difficulty.
    pub initial_threshold: u64,
    /// Difficulty for assertions between initial and graduated thresholds.
    pub reduced_difficulty: u8,
    /// Number of assertions after which PoW is exempt.
    pub graduated_threshold: u64,
    /// Trust score above which PoW is exempt.
    pub trust_exemption_score: f32,
    /// Maximum age of PoW proofs in seconds.
    pub pow_max_age: u64,
}

impl Default for AdmissionConfig {
    fn default() -> Self {
        Self {
            initial_difficulty: POW_INITIAL_DIFFICULTY,
            initial_threshold: POW_INITIAL_THRESHOLD,
            reduced_difficulty: POW_REDUCED_DIFFICULTY,
            graduated_threshold: POW_GRADUATED_THRESHOLD,
            trust_exemption_score: super::trust_tier::TRUST_POW_EXEMPTION_THRESHOLD,
            pow_max_age: POW_MAX_AGE_SECONDS,
        }
    }
}

impl AdmissionConfig {
    /// Compute the required difficulty for an agent.
    ///
    /// # Arguments
    /// * `assertion_count` - Number of assertions the agent has made
    /// * `trust_score` - Agent's current trust score
    ///
    /// # Returns
    /// Required difficulty in bits (0 = exempt).
    #[must_use]
    pub fn compute_difficulty(&self, assertion_count: u64, trust_score: f32) -> u8 {
        // Trust-based exemption
        if trust_score >= self.trust_exemption_score {
            return 0;
        }

        // Assertion-count-based graduation
        if assertion_count >= self.graduated_threshold {
            return 0;
        }

        if assertion_count >= self.initial_threshold {
            return self.reduced_difficulty;
        }

        self.initial_difficulty
    }
}

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

    #[test]
    fn test_leading_zeros_all_zero() {
        let hash = blake3::Hash::from([0u8; 32]);
        assert_eq!(PowProof::leading_zeros(&hash), 255); // Saturates at 255
    }

    #[test]
    fn test_leading_zeros_first_byte_nonzero() {
        let mut bytes = [0u8; 32];
        bytes[0] = 0x80; // 10000000 binary = 0 leading zeros in first byte
        let hash = blake3::Hash::from(bytes);
        assert_eq!(PowProof::leading_zeros(&hash), 0);

        bytes[0] = 0x40; // 01000000 binary = 1 leading zero
        let hash = blake3::Hash::from(bytes);
        assert_eq!(PowProof::leading_zeros(&hash), 1);

        bytes[0] = 0x01; // 00000001 binary = 7 leading zeros
        let hash = blake3::Hash::from(bytes);
        assert_eq!(PowProof::leading_zeros(&hash), 7);
    }

    #[test]
    fn test_leading_zeros_second_byte() {
        let mut bytes = [0u8; 32];
        bytes[0] = 0x00;
        bytes[1] = 0x80;
        let hash = blake3::Hash::from(bytes);
        assert_eq!(PowProof::leading_zeros(&hash), 8); // 8 from first byte + 0 from second
    }

    #[test]
    fn test_verify_expired_timestamp() {
        let agent_id = [0u8; 32];
        let proof = PowProof::new(0, agent_id, 1000);
        let result = proof.verify(0, 300, 2000); // 1000 seconds old, max 300

        assert!(matches!(
            result,
            Err(PowError::TimestampExpired { max_age: 300, actual_age: 1000 })
        ));
    }

    #[test]
    fn test_verify_future_timestamp() {
        let agent_id = [0u8; 32];
        let proof = PowProof::new(0, agent_id, 2000);
        let result = proof.verify(0, 300, 1000); // Proof claims timestamp 2000, server at 1000

        assert!(matches!(
            result,
            Err(PowError::TimestampInFuture { server_time: 1000, proof_time: 2000 })
        ));
    }

    #[test]
    fn test_verify_difficulty_zero_passes() {
        let agent_id = [0u8; 32];
        let now = 1700000000u64;
        let proof = PowProof::new(12345, agent_id, now);

        let result = proof.verify(0, POW_MAX_AGE_SECONDS, now);
        assert!(result.is_ok());
    }

    #[test]
    fn test_verify_insufficient_difficulty() {
        let agent_id = [0u8; 32];
        let now = 1700000000u64;
        // Random nonce unlikely to have 16 leading zeros
        let proof = PowProof::new(12345, agent_id, now);

        let result = proof.verify(16, POW_MAX_AGE_SECONDS, now);
        assert!(matches!(result, Err(PowError::InsufficientDifficulty { required: 16, .. })));
    }

    #[test]
    fn test_solve_difficulty_zero() {
        let agent_id = [1u8; 32];
        let timestamp = 1700000000u64;
        let proof = PowProof::solve(agent_id, timestamp, 0);

        assert_eq!(proof.nonce, 0);
        assert_eq!(proof.agent_id, agent_id);
        assert_eq!(proof.timestamp, timestamp);
    }

    #[test]
    fn test_solve_low_difficulty() {
        let agent_id = [2u8; 32];
        let timestamp = 1700000000u64;
        let proof = PowProof::solve(agent_id, timestamp, 4);

        // Verify the solution works
        let result = proof.verify(4, POW_MAX_AGE_SECONDS, timestamp);
        assert!(result.is_ok());
    }

    #[test]
    fn test_admission_config_default() {
        let config = AdmissionConfig::default();

        assert_eq!(config.initial_difficulty, 16);
        assert_eq!(config.initial_threshold, 10);
        assert_eq!(config.reduced_difficulty, 1);
        assert_eq!(config.graduated_threshold, 50);
        assert!((config.trust_exemption_score - 0.6).abs() < f32::EPSILON);
        assert_eq!(config.pow_max_age, 300);
    }

    #[test]
    fn test_compute_difficulty_by_assertion_count() {
        let config = AdmissionConfig::default();

        // First 10: high difficulty
        assert_eq!(config.compute_difficulty(0, 0.3), 16);
        assert_eq!(config.compute_difficulty(5, 0.3), 16);
        assert_eq!(config.compute_difficulty(9, 0.3), 16);

        // 10-49: reduced difficulty
        assert_eq!(config.compute_difficulty(10, 0.3), 1);
        assert_eq!(config.compute_difficulty(25, 0.3), 1);
        assert_eq!(config.compute_difficulty(49, 0.3), 1);

        // 50+: exempt
        assert_eq!(config.compute_difficulty(50, 0.3), 0);
        assert_eq!(config.compute_difficulty(100, 0.3), 0);
    }

    #[test]
    fn test_compute_difficulty_by_trust() {
        let config = AdmissionConfig::default();

        // Low trust, few assertions: high difficulty
        assert_eq!(config.compute_difficulty(0, 0.3), 16);

        // High trust: exempt regardless of assertion count
        assert_eq!(config.compute_difficulty(0, 0.6), 0);
        assert_eq!(config.compute_difficulty(5, 0.7), 0);
        assert_eq!(config.compute_difficulty(100, 0.9), 0);
    }

    #[test]
    fn test_hash_consistency() {
        let proof1 = PowProof::new(123, [0xAB; 32], 1700000000);
        let proof2 = PowProof::new(123, [0xAB; 32], 1700000000);

        assert_eq!(proof1.compute_hash(), proof2.compute_hash());
    }

    #[test]
    fn test_hash_changes_with_nonce() {
        let proof1 = PowProof::new(123, [0xAB; 32], 1700000000);
        let proof2 = PowProof::new(124, [0xAB; 32], 1700000000);

        assert_ne!(proof1.compute_hash(), proof2.compute_hash());
    }

    #[test]
    fn test_hash_changes_with_agent_id() {
        let proof1 = PowProof::new(123, [0xAB; 32], 1700000000);
        let proof2 = PowProof::new(123, [0xCD; 32], 1700000000);

        assert_ne!(proof1.compute_hash(), proof2.compute_hash());
    }

    #[test]
    fn test_hash_changes_with_timestamp() {
        let proof1 = PowProof::new(123, [0xAB; 32], 1700000000);
        let proof2 = PowProof::new(123, [0xAB; 32], 1700000001);

        assert_ne!(proof1.compute_hash(), proof2.compute_hash());
    }
}
