//! Phase 3 Advanced Features Example
//!
//! This example demonstrates the complete Phase 3 implementation including:
//! - Advanced 0-RTT optimization with Raw Public Keys
//! - Distributed peer discovery with certificate type awareness
//! - Enterprise certificate management with HSM support
//! - Performance optimization with ML-based tuning

use std::{
    net::{SocketAddr, IpAddr, Ipv4Addr},
    sync::Arc,
    time::Duration,
};

use tokio::time::{timeout, interval};
use tracing::{info, warn, error, Level};
use tracing_subscriber;

use ant_quic::{
    crypto::{
        // Phase 1 & 2 modules
        raw_public_keys::{RawPublicKeyConfigBuilder, utils},
        tls_extensions::{CertificateType, CertificateTypePreferences, NegotiationResult},
        certificate_negotiation::NegotiationConfig,
        quinn_integration::CertTypeQuicEndpointBuilder,
        bootstrap_support::{CertificateTypeCapabilities, BootstrapRegistryConfig},
        performance_monitoring::{
            CertTypePerformanceMonitor, PerformanceThresholds,
            ProductionCertTypeSystem, ProductionHardeningConfig,
        },
        
        // Phase 3 modules
        zero_rtt_rpk::{ZeroRttRpkConfig, ZeroRttRpkManager, RpkSessionTicket},
        peer_discovery::{CertTypeDht, DhtConfig, CertTypeGossip, GossipConfig, PeerAnnouncement},
        enterprise_cert_mgmt::{
            EnterpriseConfig, HsmConfig, HsmProvider, RotationPolicy,
            AuditConfig, BackupConfig,
            AccessControlConfig, EnterpriseCertManager,
        },
        performance_optimization::{
            OptimizationConfig, PerformanceOptimizer, PerformanceSample,
        },
    },
    nat_traversal_api::{EndpointRole, PeerId},
};

/// Advanced P2P node with all Phase 3 features
struct AdvancedP2PNode {
    /// Node identity
    peer_id: PeerId,
    /// Local keypair
    keypair: (ed25519_dalek::SigningKey, ed25519_dalek::VerifyingKey),
    /// 0-RTT manager
    zero_rtt_manager: Arc<ZeroRttRpkManager>,
    /// DHT for peer discovery
    dht: Arc<CertTypeDht>,
    /// Gossip protocol
    gossip: Arc<CertTypeGossip>,
    /// Enterprise certificate manager
    cert_manager: Option<Arc<EnterpriseCertManager>>,
    /// Performance optimizer
    optimizer: Arc<PerformanceOptimizer>,
    /// Production system (from Phase 2)
    production_system: Arc<ProductionCertTypeSystem>,
}

impl AdvancedP2PNode {
    /// Create a new advanced P2P node
    async fn new(
        bind_addr: SocketAddr,
        enable_enterprise: bool,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        info!("Creating advanced P2P node with Phase 3 features");

        // Generate identity
        let (private_key, public_key) = utils::generate_ed25519_keypair();
        let peer_id = PeerId(utils::public_key_to_bytes(&public_key));

        // Configure 0-RTT
        let zero_rtt_config = ZeroRttRpkConfig {
            enable_client: true,
            enable_server: true,
            max_early_data_size: 16384,
            ticket_lifetime: 7200,
            max_tickets: 10,
            enable_anti_replay: true,
        };
        let zero_rtt_manager = Arc::new(ZeroRttRpkManager::new(zero_rtt_config));

        // Configure DHT
        let dht_config = DhtConfig {
            bucket_size: 20,
            replication_factor: 20,
            alpha: 3,
            announce_interval: Duration::from_secs(300),
            cleanup_interval: Duration::from_secs(600),
            request_timeout: Duration::from_secs(10),
        };

        let local_capabilities = CertificateTypeCapabilities::from_preferences(
            &CertificateTypePreferences::prefer_raw_public_key(),
            EndpointRole::Client,
        );

        let (dht, mut dht_events) = CertTypeDht::new(peer_id, local_capabilities.clone(), dht_config);
        let dht = Arc::new(dht);

        // Configure gossip
        let gossip_config = GossipConfig {
            gossip_interval: Duration::from_secs(10),
            fanout: 3,
            max_peers: 1000,
            heartbeat_timeout: Duration::from_secs(60),
        };

        let (gossip, mut gossip_events) = CertTypeGossip::new(
            peer_id,
            local_capabilities,
            gossip_config,
        );
        let gossip = Arc::new(gossip);

        // Configure enterprise certificate management (if enabled)
        let cert_manager = if enable_enterprise {
            let enterprise_config = create_enterprise_config();
            let (manager, mut cert_events) = EnterpriseCertManager::new(enterprise_config).await?;
            
            // Spawn event handler
            tokio::spawn(async move {
                while let Ok(event) = cert_events.recv().await {
                    info!("Certificate management event: {:?}", event);
                }
            });
            
            Some(Arc::new(manager))
        } else {
            None
        };

        // Configure performance optimization
        let optimization_config = OptimizationConfig {
            enable_memory_pools: true,
            enable_lock_free: true,
            enable_simd_crypto: true,
            enable_ml_optimization: true,
            ..Default::default()
        };
        let optimizer = Arc::new(PerformanceOptimizer::new(optimization_config));

        // Production system from Phase 2
        let production_config = ProductionHardeningConfig::default();
        let production_system = Arc::new(ProductionCertTypeSystem::new(production_config));

        // Spawn DHT event handler
        let dht_clone = dht.clone();
        tokio::spawn(async move {
            while let Some(event) = dht_events.recv().await {
                info!("DHT event: {:?}", event);
                // Handle DHT events (peer discovery, etc.)
            }
        });

        // Spawn gossip event handler
        tokio::spawn(async move {
            while let Ok(event) = gossip_events.recv().await {
                info!("Gossip event: {:?}", event);
                // Handle gossip events
            }
        });

        info!("Advanced P2P node created: peer_id={:?}", peer_id);

        Ok(Self {
            peer_id,
            keypair: (private_key, public_key),
            zero_rtt_manager,
            dht,
            gossip,
            cert_manager,
            optimizer,
            production_system,
        })
    }

    /// Demonstrate 0-RTT session resumption
    async fn demonstrate_zero_rtt(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("=== Demonstrating 0-RTT Session Resumption ===");

        // Create a session ticket
        let peer_id = PeerId([2; 32]); // Remote peer
        let (_, remote_public_key) = utils::generate_ed25519_keypair();
        
        let cert_types = NegotiationResult::new(
            CertificateType::RawPublicKey,
            CertificateType::RawPublicKey,
        );

        let resumption_secret = vec![0u8; 32]; // Would be derived from TLS
        
        let ticket = self.zero_rtt_manager.create_session_ticket(
            peer_id,
            &remote_public_key,
            cert_types,
            Some(b"h3".to_vec()),
            resumption_secret,
            Some("example.com".to_string()),
        )?;

        info!("Created 0-RTT session ticket: session_id={:?}, age_add={}", 
              ticket.session_id, ticket.age_add);

        // Find resumption ticket
        if let Some((found_ticket, protection)) = self.zero_rtt_manager
            .find_resumption_ticket(&peer_id, Some("example.com"))
        {
            info!("Found resumption ticket for 0-RTT");
            
            // Derive early data keys
            let early_keys = protection.derive_early_data_keys()?;
            info!("Derived early data keys: client={} bytes, server={} bytes",
                  early_keys.client_key.len(), early_keys.server_key.len());
        }

        let stats = self.zero_rtt_manager.stats();
        info!("0-RTT stats: {} cached tickets, {} active protections",
              stats.cache_stats.total_tickets, stats.active_protections);

        Ok(())
    }

    /// Demonstrate distributed peer discovery
    async fn demonstrate_peer_discovery(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("=== Demonstrating Distributed Peer Discovery ===");

        // Announce ourselves to the DHT
        self.dht.announce(vec!["storage".to_string(), "compute".to_string()]).await?;
        info!("Announced to DHT with services: storage, compute");

        // Create peer announcements for testing
        for i in 0..5 {
            let peer_id = PeerId([i; 32]);
            let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, i)), 8080);
            
            let capabilities = if i % 2 == 0 {
                CertificateTypeCapabilities::from_preferences(
                    &CertificateTypePreferences::raw_public_key_only(),
                    EndpointRole::Client,
                )
            } else {
                CertificateTypeCapabilities::from_preferences(
                    &CertificateTypePreferences::prefer_raw_public_key(),
                    EndpointRole::Client,
                )
            };

            let announcement = PeerAnnouncement::new(
                peer_id,
                vec![addr],
                capabilities,
                vec!["storage".to_string()],
                3600,
            );

            self.dht.store_announcement(announcement)?;
        }

        // Find peers supporting Raw Public Keys
        let rpk_peers = self.dht.lookup_by_cert_type(CertificateType::RawPublicKey).await?;
        info!("Found {} peers supporting Raw Public Keys", rpk_peers.len());

        // Find compatible peers
        let preferences = CertificateTypePreferences::raw_public_key_only();
        let compatible = self.dht.find_compatible_peers(&preferences).await?;
        info!("Found {} compatible peers for RPK-only preferences", compatible.len());

        // Get network view from gossip
        let network_view = self.gossip.get_network_view();
        info!("Gossip network view: {} known peers", network_view.len());

        Ok(())
    }

    /// Demonstrate enterprise certificate management
    async fn demonstrate_enterprise_features(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("=== Demonstrating Enterprise Certificate Management ===");

        if let Some(cert_manager) = &self.cert_manager {
            // Generate a new Raw Public Key
            let peer_id = PeerId([100; 32]);
            let public_key = cert_manager.generate_rpk(peer_id).await?;
            info!("Generated enterprise-managed RPK: {} bytes", public_key.as_bytes().len());

            // Compliance check removed
        } else {
            info!("Enterprise features not enabled");
        }

        Ok(())
    }

    /// Demonstrate performance optimization
    async fn demonstrate_performance_optimization(&self) -> Result<(), Box<dyn std::error::Error>> {
        info!("=== Demonstrating Performance Optimization ===");

        // Allocate memory from pools
        let mut small_alloc = self.optimizer.allocate(128);
        small_alloc.get_mut().resize(128, 0);
        info!("Allocated 128 bytes from memory pool");

        // Cache negotiation results
        let result = NegotiationResult::new(
            CertificateType::RawPublicKey,
            CertificateType::RawPublicKey,
        );
        
        self.optimizer.cache_negotiation(12345, result.clone());
        
        if let Some(cached) = self.optimizer.lookup_negotiation(12345) {
            info!("Retrieved cached negotiation result");
        }

        // Record performance samples
        for i in 0..10 {
            let sample = PerformanceSample {
                timestamp: std::time::Instant::now(),
                negotiation_time: Duration::from_millis(5 + i),
                cache_hit: i % 2 == 0,
                used_rpk: true,
            };
            
            self.optimizer.record_performance(sample);
        }

        // Get optimization recommendations
        let recommendations = self.optimizer.get_recommendations();
        info!("ML-based recommendations:");
        info!("  - Cache size: {}", recommendations.recommended_cache_size);
        info!("  - Pool size: {}", recommendations.recommended_pool_size);
        info!("  - Use SIMD: {}", recommendations.use_simd);
        info!("  - Prefetch: {}", recommendations.prefetch_hint);

        // Get performance stats
        let stats = self.optimizer.get_stats();
        info!("Performance stats:");
        info!("  - Small pool hit rate: {:.2}%", stats.small_pool_stats.hit_rate() * 100.0);
        info!("  - Cache size: {}", stats.cache_size);

        Ok(())
    }

    /// Run periodic maintenance tasks
    async fn run_maintenance(&self) {
        let mut interval = interval(Duration::from_secs(60));
        
        loop {
            interval.tick().await;
            
            // Clean up 0-RTT sessions
            self.zero_rtt_manager.cleanup();
            
            // Clean up DHT
            self.dht.cleanup();
            
            // Get health report
            let health = self.production_system.health_report();
            info!("System health: {:?}", health.status);
        }
    }
}

/// Create enterprise configuration
fn create_enterprise_config() -> EnterpriseConfig {
    use std::collections::HashMap;
    use std::path::PathBuf;
    
    EnterpriseConfig {
        hsm: Some(HsmConfig {
            provider: HsmProvider::SoftwareEmulation,
            connection: HashMap::new(),
            key_policy: ant_quic::crypto::enterprise_cert_mgmt::HsmKeyPolicy {
                store_rpk: true,
                store_x509: true,
                allowed_algorithms: vec!["Ed25519".to_string()],
                min_key_size: 256,
                require_attestation: false,
            },
            performance: ant_quic::crypto::enterprise_cert_mgmt::HsmPerformance {
                pool_size: 10,
                timeout: Duration::from_secs(5),
                cache_public_keys: true,
                enable_batching: true,
            },
        }),
        rotation_policy: RotationPolicy {
            auto_rotate: true,
            interval: Duration::from_secs(86400 * 90), // 90 days
            grace_period: Duration::from_secs(86400 * 7), // 7 days
            max_age: Duration::from_secs(86400 * 365), // 1 year
            triggers: vec![],
        },
        audit: AuditConfig {
            enabled: true,
            destination: ant_quic::crypto::enterprise_cert_mgmt::AuditDestination::File { 
                path: PathBuf::from("/tmp/cert-audit.log") 
            },
            events: vec![],
            detailed_context: true,
        },
        backup: BackupConfig {
            enabled: false,
            destinations: vec![],
            encryption: ant_quic::crypto::enterprise_cert_mgmt::BackupEncryption {
                algorithm: "AES-256-GCM".to_string(),
                kdf: "PBKDF2".to_string(),
                master_key_location: "memory".to_string(),
            },
            schedule: ant_quic::crypto::enterprise_cert_mgmt::BackupSchedule::Daily { hour: 2, minute: 0 },
        },
        access_control: AccessControlConfig {
            rbac_enabled: true,
            roles: HashMap::new(),
            mfa_required: false,
            session_timeout: Duration::from_secs(3600),
        },
    }
}

/// Main demonstration
async fn run_demonstration() -> Result<(), Box<dyn std::error::Error>> {
    info!("Starting Phase 3 Advanced Features Demonstration");

    // Create nodes
    let node1_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9001);
    let node1 = AdvancedP2PNode::new(node1_addr, true).await?;

    let node2_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9002);
    let node2 = AdvancedP2PNode::new(node2_addr, false).await?;

    // Run demonstrations
    node1.demonstrate_zero_rtt().await?;
    tokio::time::sleep(Duration::from_secs(1)).await;

    node1.demonstrate_peer_discovery().await?;
    tokio::time::sleep(Duration::from_secs(1)).await;

    node1.demonstrate_enterprise_features().await?;
    tokio::time::sleep(Duration::from_secs(1)).await;

    node1.demonstrate_performance_optimization().await?;

    // Start maintenance task
    let node1_clone = Arc::new(node1);
    tokio::spawn(async move {
        node1_clone.run_maintenance().await;
    });

    info!("Phase 3 demonstration completed successfully!");
    info!("The system is now running with:");
    info!("  - 0-RTT session resumption for Raw Public Keys");
    info!("  - Distributed peer discovery with DHT and gossip");
    info!("  - Enterprise certificate management with HSM support");
    info!("  - ML-based performance optimization");

    // Keep running for monitoring
    tokio::time::sleep(Duration::from_secs(5)).await;

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .init();

    info!("=================================================");
    info!("ANT-QUIC Phase 3: Advanced Features Demonstration");
    info!("=================================================");

    if let Err(e) = run_demonstration().await {
        error!("Demonstration failed: {}", e);
        std::process::exit(1);
    }

    Ok(())
}