//! Post-Quantum Algorithm Negotiation Protocol
//!
//! This module provides a protocol for nodes to negotiate which post-quantum
//! signature algorithm to use for communication. Nodes advertise their
//! capabilities and preferences, and the protocol selects the best common
//! algorithm based on security, performance, and bandwidth constraints.
//!
//! # Supported Algorithms
//!
//! | Algorithm | NIST | Sig Size | Use Case |
//! |-----------|------|----------|----------|
//! | Dilithium3 (ML-DSA-65) | FIPS 204 | ~3.3KB | Default, balanced |
//! | SPHINCS+ (SLH-DSA) | FIPS 205 | ~7.8KB | Conservative, hash-based |
//! | FALCON (FN-DSA) | FIPS 206 | ~0.7KB | Compact, mobile-friendly |
//!
//! # Negotiation Flow
//!
//! ```text
//! Node A                          Node B
//!   │                                │
//!   ├──── Capabilities ─────────────►│
//!   │     (supported algorithms,     │
//!   │      preferences, version)     │
//!   │                                │
//!   │◄─── Capabilities ──────────────┤
//!   │                                │
//!   ├──── Selection ────────────────►│
//!   │     (chosen algorithm,         │
//!   │      session parameters)       │
//!   │                                │
//!   │◄─── Acknowledgment ────────────┤
//!   │                                │
//!   │ ═══ Secure Channel ══════════  │
//!   │     (using negotiated algo)    │
//! ```
//!
//! # Selection Priority
//!
//! 1. Security: Prefer algorithms with higher security margins
//! 2. Compatibility: Must be supported by both nodes
//! 3. Bandwidth: Consider signature size for constrained networks
//! 4. Performance: Prefer faster algorithms when security is equal

use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use thiserror::Error;

/// Post-quantum signature algorithms supported by Synor
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum PqAlgorithm {
    /// Dilithium3 (ML-DSA-65) - NIST FIPS 204
    /// Default choice, balanced security and performance
    Dilithium3,

    /// SPHINCS+-SHAKE-128s (SLH-DSA) - NIST FIPS 205
    /// Hash-based, conservative security, larger signatures
    SphincsShake128s,

    /// SPHINCS+-SHAKE-192s (SLH-DSA) - NIST FIPS 205
    /// Higher security margin, larger signatures
    SphincsShake192s,

    /// SPHINCS+-SHAKE-256s (SLH-DSA) - NIST FIPS 205
    /// Maximum security, largest signatures
    SphincsShake256s,

    /// FALCON-512 (FN-DSA) - NIST FIPS 206
    /// Compact signatures, ideal for mobile/IoT
    Falcon512,

    /// FALCON-1024 (FN-DSA) - NIST FIPS 206
    /// Higher security with compact signatures
    Falcon1024,
}

impl PqAlgorithm {
    /// Get NIST security level (1-5)
    pub fn security_level(&self) -> u8 {
        match self {
            Self::Dilithium3 => 3,
            Self::SphincsShake128s => 1,
            Self::SphincsShake192s => 3,
            Self::SphincsShake256s => 5,
            Self::Falcon512 => 1,
            Self::Falcon1024 => 5,
        }
    }

    /// Get approximate signature size in bytes
    pub fn signature_size(&self) -> usize {
        match self {
            Self::Dilithium3 => 3293,
            Self::SphincsShake128s => 7856,
            Self::SphincsShake192s => 16224,
            Self::SphincsShake256s => 29792,
            Self::Falcon512 => 690,
            Self::Falcon1024 => 1330,
        }
    }

    /// Get approximate public key size in bytes
    pub fn public_key_size(&self) -> usize {
        match self {
            Self::Dilithium3 => 1952,
            Self::SphincsShake128s => 32,
            Self::SphincsShake192s => 48,
            Self::SphincsShake256s => 64,
            Self::Falcon512 => 897,
            Self::Falcon1024 => 1793,
        }
    }

    /// Get the NIST FIPS standard number
    pub fn fips_standard(&self) -> &'static str {
        match self {
            Self::Dilithium3 => "FIPS 204 (ML-DSA)",
            Self::SphincsShake128s | Self::SphincsShake192s | Self::SphincsShake256s => {
                "FIPS 205 (SLH-DSA)"
            }
            Self::Falcon512 | Self::Falcon1024 => "FIPS 206 (FN-DSA)",
        }
    }

    /// Check if this is a hash-based algorithm (SPHINCS+)
    pub fn is_hash_based(&self) -> bool {
        matches!(
            self,
            Self::SphincsShake128s | Self::SphincsShake192s | Self::SphincsShake256s
        )
    }

    /// Check if this is a lattice-based algorithm
    pub fn is_lattice_based(&self) -> bool {
        matches!(self, Self::Dilithium3 | Self::Falcon512 | Self::Falcon1024)
    }

    /// Get human-readable name
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::Dilithium3 => "Dilithium3 (ML-DSA-65)",
            Self::SphincsShake128s => "SPHINCS+-SHAKE-128s",
            Self::SphincsShake192s => "SPHINCS+-SHAKE-192s",
            Self::SphincsShake256s => "SPHINCS+-SHAKE-256s",
            Self::Falcon512 => "FALCON-512",
            Self::Falcon1024 => "FALCON-1024",
        }
    }

    /// Default priority order (higher = more preferred)
    fn default_priority(&self) -> u8 {
        match self {
            Self::Dilithium3 => 100,      // Default, well-balanced
            Self::Falcon1024 => 90,       // High security, compact
            Self::Falcon512 => 85,        // Compact, mobile-friendly
            Self::SphincsShake192s => 70, // Conservative backup
            Self::SphincsShake256s => 60, // Maximum security
            Self::SphincsShake128s => 50, // Basic SPHINCS+
        }
    }

    /// All supported algorithms
    pub fn all() -> &'static [PqAlgorithm] {
        &[
            Self::Dilithium3,
            Self::SphincsShake128s,
            Self::SphincsShake192s,
            Self::SphincsShake256s,
            Self::Falcon512,
            Self::Falcon1024,
        ]
    }
}

impl std::fmt::Display for PqAlgorithm {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.display_name())
    }
}

/// Node's algorithm capabilities and preferences
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AlgorithmCapabilities {
    /// Protocol version
    pub version: u32,
    /// Node identifier (public key hash or peer ID)
    pub node_id: [u8; 32],
    /// Supported algorithms with their priorities (higher = more preferred)
    pub supported: HashMap<PqAlgorithm, u8>,
    /// Minimum acceptable security level (NIST 1-5)
    pub min_security_level: u8,
    /// Maximum acceptable signature size (0 = unlimited)
    pub max_signature_size: usize,
    /// Preferred algorithm family (lattice, hash, any)
    pub preferred_family: AlgorithmFamily,
    /// Timestamp for freshness
    pub timestamp: u64,
    /// Optional extensions for future compatibility
    pub extensions: HashMap<String, Vec<u8>>,
}

impl AlgorithmCapabilities {
    /// Create capabilities with default preferences
    pub fn default_for_node(node_id: [u8; 32]) -> Self {
        let mut supported = HashMap::new();
        for algo in PqAlgorithm::all() {
            supported.insert(*algo, algo.default_priority());
        }

        Self {
            version: 1,
            node_id,
            supported,
            min_security_level: 1,
            max_signature_size: 0,
            preferred_family: AlgorithmFamily::Any,
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
            extensions: HashMap::new(),
        }
    }

    /// Create capabilities for bandwidth-constrained environments
    pub fn bandwidth_optimized(node_id: [u8; 32]) -> Self {
        let mut caps = Self::default_for_node(node_id);
        caps.max_signature_size = 5000; // Exclude large SPHINCS+ variants
        caps.supported.insert(PqAlgorithm::Falcon512, 100);
        caps.supported.insert(PqAlgorithm::Falcon1024, 95);
        caps.supported.insert(PqAlgorithm::Dilithium3, 80);
        caps
    }

    /// Create capabilities for maximum security
    pub fn security_optimized(node_id: [u8; 32]) -> Self {
        let mut caps = Self::default_for_node(node_id);
        caps.min_security_level = 3;
        caps.supported.insert(PqAlgorithm::SphincsShake256s, 100);
        caps.supported.insert(PqAlgorithm::Falcon1024, 95);
        caps.supported.insert(PqAlgorithm::SphincsShake192s, 90);
        caps.supported.insert(PqAlgorithm::Dilithium3, 85);
        caps
    }

    /// Add support for an algorithm with priority
    pub fn add_algorithm(&mut self, algo: PqAlgorithm, priority: u8) {
        self.supported.insert(algo, priority);
    }

    /// Remove support for an algorithm
    pub fn remove_algorithm(&mut self, algo: PqAlgorithm) {
        self.supported.remove(&algo);
    }

    /// Check if an algorithm is supported
    pub fn supports(&self, algo: PqAlgorithm) -> bool {
        self.supported.contains_key(&algo)
    }

    /// Get priority for an algorithm (0 if not supported)
    pub fn priority(&self, algo: PqAlgorithm) -> u8 {
        self.supported.get(&algo).copied().unwrap_or(0)
    }

    /// Encode capabilities to bytes for transmission
    pub fn encode(&self) -> Vec<u8> {
        serde_json::to_vec(self).unwrap_or_default()
    }

    /// Decode capabilities from bytes
    pub fn decode(data: &[u8]) -> Result<Self, NegotiationError> {
        serde_json::from_slice(data)
            .map_err(|e| NegotiationError::InvalidCapabilities(e.to_string()))
    }
}

/// Algorithm family preference
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum AlgorithmFamily {
    /// Lattice-based (Dilithium, FALCON)
    Lattice,
    /// Hash-based (SPHINCS+)
    HashBased,
    /// No preference
    #[default]
    Any,
}

/// Result of algorithm negotiation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NegotiationResult {
    /// Selected algorithm
    pub algorithm: PqAlgorithm,
    /// Combined score (higher = better match)
    pub score: u32,
    /// Whether both nodes preferred this algorithm
    pub mutual_preference: bool,
    /// Session parameters
    pub session_params: SessionParams,
}

/// Session parameters established during negotiation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionParams {
    /// Session identifier
    pub session_id: [u8; 32],
    /// Negotiation timestamp
    pub negotiated_at: u64,
    /// Session expiry (0 = no expiry)
    pub expires_at: u64,
    /// Allow algorithm renegotiation mid-session
    pub allow_renegotiation: bool,
    /// Fallback algorithm if primary fails
    pub fallback: Option<PqAlgorithm>,
}

impl Default for SessionParams {
    fn default() -> Self {
        use std::time::{SystemTime, UNIX_EPOCH};

        let mut session_id = [0u8; 32];
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        session_id[..16].copy_from_slice(&now.to_le_bytes());

        Self {
            session_id,
            negotiated_at: now as u64 / 1_000_000_000,
            expires_at: 0,
            allow_renegotiation: true,
            fallback: Some(PqAlgorithm::Dilithium3),
        }
    }
}

/// Algorithm negotiator
pub struct AlgorithmNegotiator {
    /// Local node capabilities
    local_caps: AlgorithmCapabilities,
    /// Negotiation policy
    policy: NegotiationPolicy,
}

impl AlgorithmNegotiator {
    /// Create a new negotiator with local capabilities
    pub fn new(local_caps: AlgorithmCapabilities) -> Self {
        Self {
            local_caps,
            policy: NegotiationPolicy::default(),
        }
    }

    /// Create with custom policy
    pub fn with_policy(local_caps: AlgorithmCapabilities, policy: NegotiationPolicy) -> Self {
        Self { local_caps, policy }
    }

    /// Get local capabilities
    pub fn capabilities(&self) -> &AlgorithmCapabilities {
        &self.local_caps
    }

    /// Negotiate algorithm with remote peer
    pub fn negotiate(
        &self,
        remote_caps: &AlgorithmCapabilities,
    ) -> Result<NegotiationResult, NegotiationError> {
        // Find common algorithms
        let local_algos: HashSet<_> = self.local_caps.supported.keys().collect();
        let remote_algos: HashSet<_> = remote_caps.supported.keys().collect();
        let common: Vec<_> = local_algos.intersection(&remote_algos).copied().collect();

        if common.is_empty() {
            return Err(NegotiationError::NoCommonAlgorithm);
        }

        // Filter by constraints
        let valid_algos: Vec<_> = common
            .into_iter()
            .filter(|algo| {
                // Check security level
                let meets_local_security =
                    algo.security_level() >= self.local_caps.min_security_level;
                let meets_remote_security = algo.security_level() >= remote_caps.min_security_level;

                // Check signature size
                let local_size_ok = self.local_caps.max_signature_size == 0
                    || algo.signature_size() <= self.local_caps.max_signature_size;
                let remote_size_ok = remote_caps.max_signature_size == 0
                    || algo.signature_size() <= remote_caps.max_signature_size;

                // Check family preference
                let family_ok = self.matches_family_preference(algo);

                meets_local_security
                    && meets_remote_security
                    && local_size_ok
                    && remote_size_ok
                    && family_ok
            })
            .collect();

        if valid_algos.is_empty() {
            return Err(NegotiationError::ConstraintsMismatch);
        }

        // Score and select best algorithm
        let mut best_algo = valid_algos[0];
        let mut best_score = 0u32;

        for algo in valid_algos {
            let local_prio = self.local_caps.priority(*algo) as u32;
            let remote_prio = remote_caps.priority(*algo) as u32;

            let score = match self.policy.scoring {
                ScoringStrategy::LocalPreference => local_prio * 2 + remote_prio,
                ScoringStrategy::RemotePreference => local_prio + remote_prio * 2,
                ScoringStrategy::Average => (local_prio + remote_prio) / 2,
                ScoringStrategy::Minimum => local_prio.min(remote_prio),
                ScoringStrategy::Maximum => local_prio.max(remote_prio),
            };

            // Bonus for security level
            let security_bonus = if self.policy.prioritize_security {
                algo.security_level() as u32 * 10
            } else {
                0
            };

            // Penalty for large signatures (if bandwidth-conscious)
            let size_penalty = if self.policy.minimize_bandwidth {
                (algo.signature_size() / 1000) as u32
            } else {
                0
            };

            let final_score = score + security_bonus - size_penalty;

            if final_score > best_score {
                best_score = final_score;
                best_algo = algo;
            }
        }

        // Determine mutual preference
        let local_top = self.top_preference();
        let remote_top = Self::top_preference_from_caps(remote_caps);
        let mutual_preference = local_top == Some(*best_algo) && remote_top == Some(*best_algo);

        // Create session params
        let session_params = SessionParams {
            fallback: self.select_fallback(best_algo, remote_caps),
            ..Default::default()
        };

        Ok(NegotiationResult {
            algorithm: *best_algo,
            score: best_score,
            mutual_preference,
            session_params,
        })
    }

    /// Check if algorithm matches family preference
    fn matches_family_preference(&self, algo: &PqAlgorithm) -> bool {
        match self.local_caps.preferred_family {
            AlgorithmFamily::Any => true,
            AlgorithmFamily::Lattice => algo.is_lattice_based(),
            AlgorithmFamily::HashBased => algo.is_hash_based(),
        }
    }

    /// Get local top preference
    fn top_preference(&self) -> Option<PqAlgorithm> {
        self.local_caps
            .supported
            .iter()
            .max_by_key(|(_, prio)| *prio)
            .map(|(algo, _)| *algo)
    }

    /// Get top preference from capabilities
    fn top_preference_from_caps(caps: &AlgorithmCapabilities) -> Option<PqAlgorithm> {
        caps.supported
            .iter()
            .max_by_key(|(_, prio)| *prio)
            .map(|(algo, _)| *algo)
    }

    /// Select a fallback algorithm different from primary
    fn select_fallback(
        &self,
        primary: &PqAlgorithm,
        remote_caps: &AlgorithmCapabilities,
    ) -> Option<PqAlgorithm> {
        // Prefer a different algorithm family for fallback
        let primary_is_lattice = primary.is_lattice_based();

        self.local_caps
            .supported
            .keys()
            .filter(|algo| {
                **algo != *primary
                    && remote_caps.supports(**algo)
                    && algo.is_lattice_based() != primary_is_lattice
            })
            .max_by_key(|algo| self.local_caps.priority(**algo))
            .copied()
    }

    /// Quick negotiation using just algorithm names
    pub fn quick_negotiate(local: &[PqAlgorithm], remote: &[PqAlgorithm]) -> Option<PqAlgorithm> {
        // Find common algorithms and return the one with highest default priority
        let local_set: HashSet<_> = local.iter().collect();
        let remote_set: HashSet<_> = remote.iter().collect();

        local_set
            .intersection(&remote_set)
            .max_by_key(|algo| algo.default_priority())
            .copied()
            .copied()
    }
}

/// Negotiation policy configuration
#[derive(Debug, Clone)]
pub struct NegotiationPolicy {
    /// How to score algorithm preferences
    pub scoring: ScoringStrategy,
    /// Prioritize higher security levels
    pub prioritize_security: bool,
    /// Minimize bandwidth usage (prefer smaller signatures)
    pub minimize_bandwidth: bool,
    /// Require mutual top preference
    pub require_mutual: bool,
}

impl Default for NegotiationPolicy {
    fn default() -> Self {
        Self {
            scoring: ScoringStrategy::Average,
            prioritize_security: false,
            minimize_bandwidth: false,
            require_mutual: false,
        }
    }
}

/// Strategy for scoring algorithm preferences
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ScoringStrategy {
    /// Weight local preferences higher
    LocalPreference,
    /// Weight remote preferences higher
    RemotePreference,
    /// Average both preferences
    Average,
    /// Use minimum (both must prefer)
    Minimum,
    /// Use maximum (either prefers)
    Maximum,
}

/// Negotiation errors
#[derive(Debug, Error)]
pub enum NegotiationError {
    #[error("No common algorithm supported by both nodes")]
    NoCommonAlgorithm,

    #[error("Constraints cannot be satisfied (security level or size limits)")]
    ConstraintsMismatch,

    #[error("Invalid capabilities message: {0}")]
    InvalidCapabilities(String),

    #[error("Protocol version mismatch: local={local}, remote={remote}")]
    VersionMismatch { local: u32, remote: u32 },

    #[error("Negotiation timeout")]
    Timeout,

    #[error("Remote rejected negotiation: {0}")]
    Rejected(String),
}

/// Negotiation message types for the protocol
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum NegotiationMessage {
    /// Advertise capabilities
    Capabilities(AlgorithmCapabilities),

    /// Propose selected algorithm
    Selection {
        algorithm: PqAlgorithm,
        session_params: SessionParams,
    },

    /// Acknowledge selection
    Acknowledgment {
        session_id: [u8; 32],
        accepted: bool,
    },

    /// Request renegotiation
    Renegotiate { reason: String },

    /// Error message
    Error { code: u32, message: String },
}

impl NegotiationMessage {
    /// Encode message to bytes
    pub fn encode(&self) -> Vec<u8> {
        serde_json::to_vec(self).unwrap_or_default()
    }

    /// Decode message from bytes
    pub fn decode(data: &[u8]) -> Result<Self, NegotiationError> {
        serde_json::from_slice(data)
            .map_err(|e| NegotiationError::InvalidCapabilities(e.to_string()))
    }
}

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

    #[test]
    fn test_algorithm_properties() {
        assert_eq!(PqAlgorithm::Dilithium3.security_level(), 3);
        assert_eq!(PqAlgorithm::Falcon512.signature_size(), 690);
        assert!(PqAlgorithm::SphincsShake256s.is_hash_based());
        assert!(PqAlgorithm::Dilithium3.is_lattice_based());
    }

    #[test]
    fn test_capabilities_default() {
        let node_id = [0xabu8; 32];
        let caps = AlgorithmCapabilities::default_for_node(node_id);

        assert_eq!(caps.version, 1);
        assert!(caps.supports(PqAlgorithm::Dilithium3));
        assert!(caps.supports(PqAlgorithm::Falcon512));
        assert!(caps.supports(PqAlgorithm::SphincsShake128s));
    }

    #[test]
    fn test_bandwidth_optimized_caps() {
        let node_id = [0xabu8; 32];
        let caps = AlgorithmCapabilities::bandwidth_optimized(node_id);

        // FALCON should have highest priority
        assert!(caps.priority(PqAlgorithm::Falcon512) > caps.priority(PqAlgorithm::Dilithium3));
    }

    #[test]
    fn test_basic_negotiation() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        let local_caps = AlgorithmCapabilities::default_for_node(local_id);
        let remote_caps = AlgorithmCapabilities::default_for_node(remote_id);

        let negotiator = AlgorithmNegotiator::new(local_caps);
        let result = negotiator.negotiate(&remote_caps).unwrap();

        // Should select Dilithium3 as default preference
        assert_eq!(result.algorithm, PqAlgorithm::Dilithium3);
        assert!(result.mutual_preference);
    }

    #[test]
    fn test_bandwidth_constrained_negotiation() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        let local_caps = AlgorithmCapabilities::bandwidth_optimized(local_id);
        let remote_caps = AlgorithmCapabilities::bandwidth_optimized(remote_id);

        let policy = NegotiationPolicy {
            minimize_bandwidth: true,
            ..Default::default()
        };

        let negotiator = AlgorithmNegotiator::with_policy(local_caps, policy);
        let result = negotiator.negotiate(&remote_caps).unwrap();

        // Should prefer FALCON for bandwidth-constrained scenarios
        assert!(
            result.algorithm == PqAlgorithm::Falcon512
                || result.algorithm == PqAlgorithm::Falcon1024
        );
    }

    #[test]
    fn test_security_constrained_negotiation() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        let local_caps = AlgorithmCapabilities::security_optimized(local_id);
        let remote_caps = AlgorithmCapabilities::security_optimized(remote_id);

        let negotiator = AlgorithmNegotiator::new(local_caps);
        let result = negotiator.negotiate(&remote_caps).unwrap();

        // Should select high-security algorithm
        assert!(result.algorithm.security_level() >= 3);
    }

    #[test]
    fn test_no_common_algorithm() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        let mut local_caps = AlgorithmCapabilities::default_for_node(local_id);
        local_caps.supported.clear();
        local_caps.add_algorithm(PqAlgorithm::Falcon512, 100);

        let mut remote_caps = AlgorithmCapabilities::default_for_node(remote_id);
        remote_caps.supported.clear();
        remote_caps.add_algorithm(PqAlgorithm::SphincsShake256s, 100);

        let negotiator = AlgorithmNegotiator::new(local_caps);
        let result = negotiator.negotiate(&remote_caps);

        assert!(matches!(result, Err(NegotiationError::NoCommonAlgorithm)));
    }

    #[test]
    fn test_quick_negotiate() {
        let local = vec![PqAlgorithm::Dilithium3, PqAlgorithm::Falcon512];
        let remote = vec![PqAlgorithm::Falcon512, PqAlgorithm::SphincsShake128s];

        let result = AlgorithmNegotiator::quick_negotiate(&local, &remote);
        assert_eq!(result, Some(PqAlgorithm::Falcon512));
    }

    #[test]
    fn test_message_encoding() {
        let node_id = [0xaa; 32];
        let caps = AlgorithmCapabilities::default_for_node(node_id);
        let msg = NegotiationMessage::Capabilities(caps);

        let encoded = msg.encode();
        let decoded = NegotiationMessage::decode(&encoded).unwrap();

        if let NegotiationMessage::Capabilities(decoded_caps) = decoded {
            assert_eq!(decoded_caps.node_id, node_id);
        } else {
            panic!("Wrong message type");
        }
    }

    #[test]
    fn test_fallback_selection() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        let local_caps = AlgorithmCapabilities::default_for_node(local_id);
        let remote_caps = AlgorithmCapabilities::default_for_node(remote_id);

        let negotiator = AlgorithmNegotiator::new(local_caps);
        let result = negotiator.negotiate(&remote_caps).unwrap();

        // Fallback should be different from primary
        if let Some(fallback) = result.session_params.fallback {
            assert_ne!(fallback, result.algorithm);
            // Fallback should be from different family if possible
            if result.algorithm.is_lattice_based() {
                // Might be hash-based fallback
            }
        }
    }

    #[test]
    fn test_constraints_mismatch() {
        let local_id = [0xaa; 32];
        let remote_id = [0xbb; 32];

        // Local requires security level 5
        let mut local_caps = AlgorithmCapabilities::default_for_node(local_id);
        local_caps.min_security_level = 5;

        // Remote only supports low-security algorithms
        let mut remote_caps = AlgorithmCapabilities::default_for_node(remote_id);
        remote_caps.supported.clear();
        remote_caps.add_algorithm(PqAlgorithm::Falcon512, 100); // Level 1
        remote_caps.add_algorithm(PqAlgorithm::SphincsShake128s, 90); // Level 1

        let negotiator = AlgorithmNegotiator::new(local_caps);
        let result = negotiator.negotiate(&remote_caps);

        assert!(matches!(result, Err(NegotiationError::ConstraintsMismatch)));
    }
}