//! Tiered JIT Compilation.
//!
//! Implements multi-tier compilation strategy inspired by V8 and HotSpot:
//!
//! - **Tier 0 (Cold)**: Minimal optimization, fast compile
//! - **Tier 1 (Warm)**: Balanced optimization for frequently called contracts
//! - **Tier 2 (Hot)**: Maximum optimization for heavily used contracts
//!
//! # Benefits
//!
//! - Fast initial execution (no upfront compilation cost)
//! - Optimal performance for hot contracts
//! - Adaptive to workload patterns
//! - Efficient memory use through LRU eviction

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use lru::LruCache;
use parking_lot::{Mutex, RwLock};
use wasmtime::{Config, Engine, Module, OptLevel};

use crate::{ContractId, VmError, MAX_CONTRACT_SIZE};

/// Compilation tier levels.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CompilationTier {
    /// Cold - minimal optimization, fastest compile time.
    Cold = 0,
    /// Warm - balanced optimization.
    Warm = 1,
    /// Hot - maximum optimization, slowest compile time.
    Hot = 2,
}

impl CompilationTier {
    /// Maps tier to Wasmtime optimization level.
    pub fn opt_level(&self) -> OptLevel {
        match self {
            CompilationTier::Cold => OptLevel::SpeedAndSize,
            CompilationTier::Warm => OptLevel::Speed,
            CompilationTier::Hot => OptLevel::Speed, // Maximum in wasmtime
        }
    }

    /// Returns the execution threshold to reach this tier.
    pub fn threshold(&self) -> u64 {
        match self {
            CompilationTier::Cold => 0,
            CompilationTier::Warm => 10,   // After 10 executions
            CompilationTier::Hot => 100,   // After 100 executions
        }
    }
}

/// Configuration for tiered compilation.
#[derive(Clone, Debug)]
pub struct TieredConfig {
    /// Maximum cached modules per tier.
    pub max_cached_per_tier: usize,
    /// Enable background recompilation.
    pub background_recompile: bool,
    /// Threshold for warm tier (execution count).
    pub warm_threshold: u64,
    /// Threshold for hot tier (execution count).
    pub hot_threshold: u64,
    /// Time-to-live for cold cache entries.
    pub cold_ttl: Duration,
    /// Enable AOT caching to disk.
    pub aot_cache_enabled: bool,
    /// Path for AOT cache.
    pub aot_cache_path: Option<String>,
}

impl Default for TieredConfig {
    fn default() -> Self {
        Self {
            max_cached_per_tier: 256,
            background_recompile: true,
            warm_threshold: 10,
            hot_threshold: 100,
            cold_ttl: Duration::from_secs(300), // 5 minutes
            aot_cache_enabled: false,
            aot_cache_path: None,
        }
    }
}

/// Statistics for a compiled contract.
#[derive(Clone, Debug)]
pub struct ContractStats {
    /// Number of executions.
    pub execution_count: u64,
    /// Total gas consumed.
    pub total_gas: u64,
    /// Average execution time (microseconds).
    pub avg_exec_time_us: u64,
    /// Last execution timestamp.
    pub last_executed: Instant,
    /// Current compilation tier.
    pub tier: CompilationTier,
    /// Compile time (microseconds).
    pub compile_time_us: u64,
}

impl Default for ContractStats {
    fn default() -> Self {
        Self {
            execution_count: 0,
            total_gas: 0,
            avg_exec_time_us: 0,
            last_executed: Instant::now(),
            tier: CompilationTier::Cold,
            compile_time_us: 0,
        }
    }
}

/// A compiled module with tier information.
pub struct TieredModule {
    /// Contract ID.
    pub id: ContractId,
    /// Original bytecode.
    pub bytecode: Vec<u8>,
    /// Compiled module.
    pub module: Module,
    /// Compilation tier.
    pub tier: CompilationTier,
    /// Compile time.
    pub compile_time: Duration,
}

impl TieredModule {
    /// Creates a new tiered module.
    pub fn new(
        id: ContractId,
        bytecode: Vec<u8>,
        module: Module,
        tier: CompilationTier,
        compile_time: Duration,
    ) -> Self {
        Self {
            id,
            bytecode,
            module,
            tier,
            compile_time,
        }
    }
}

/// Engine cache for different optimization tiers.
struct TierEngines {
    cold: Engine,
    warm: Engine,
    hot: Engine,
}

impl TierEngines {
    fn new() -> Result<Self, VmError> {
        Ok(Self {
            cold: Self::create_engine(CompilationTier::Cold)?,
            warm: Self::create_engine(CompilationTier::Warm)?,
            hot: Self::create_engine(CompilationTier::Hot)?,
        })
    }

    fn create_engine(tier: CompilationTier) -> Result<Engine, VmError> {
        let mut config = Config::new();
        config.consume_fuel(true);
        config.wasm_bulk_memory(true);
        config.wasm_multi_value(true);
        config.wasm_reference_types(true);
        config.cranelift_opt_level(tier.opt_level());

        // Hot tier gets additional optimizations
        if tier == CompilationTier::Hot {
            config.parallel_compilation(true);
        }

        Engine::new(&config).map_err(|e| VmError::ExecutionError(e.to_string()))
    }

    fn get(&self, tier: CompilationTier) -> &Engine {
        match tier {
            CompilationTier::Cold => &self.cold,
            CompilationTier::Warm => &self.warm,
            CompilationTier::Hot => &self.hot,
        }
    }
}

/// The tiered JIT compiler.
pub struct TieredCompiler {
    /// Configuration.
    config: TieredConfig,
    /// Engines for each tier.
    engines: TierEngines,
    /// Cached modules by tier.
    cold_cache: Mutex<LruCache<ContractId, Arc<TieredModule>>>,
    warm_cache: Mutex<LruCache<ContractId, Arc<TieredModule>>>,
    hot_cache: RwLock<HashMap<ContractId, Arc<TieredModule>>>,
    /// Execution statistics per contract.
    stats: RwLock<HashMap<ContractId, ContractStats>>,
    /// Bytecode storage for recompilation.
    bytecode_store: RwLock<HashMap<ContractId, Vec<u8>>>,
    /// Global statistics.
    total_compilations: AtomicU64,
    total_cache_hits: AtomicU64,
    total_cache_misses: AtomicU64,
}

impl TieredCompiler {
    /// Creates a new tiered compiler.
    pub fn new() -> Result<Self, VmError> {
        Self::with_config(TieredConfig::default())
    }

    /// Creates with custom configuration.
    pub fn with_config(config: TieredConfig) -> Result<Self, VmError> {
        let max_cached = config.max_cached_per_tier;
        Ok(Self {
            config,
            engines: TierEngines::new()?,
            cold_cache: Mutex::new(LruCache::new(
                std::num::NonZeroUsize::new(max_cached).unwrap(),
            )),
            warm_cache: Mutex::new(LruCache::new(
                std::num::NonZeroUsize::new(max_cached).unwrap(),
            )),
            hot_cache: RwLock::new(HashMap::new()),
            stats: RwLock::new(HashMap::new()),
            bytecode_store: RwLock::new(HashMap::new()),
            total_compilations: AtomicU64::new(0),
            total_cache_hits: AtomicU64::new(0),
            total_cache_misses: AtomicU64::new(0),
        })
    }

    /// Gets or compiles a module at the appropriate tier.
    pub fn get_or_compile(&self, id: &ContractId, bytecode: &[u8]) -> Result<Arc<TieredModule>, VmError> {
        // Check hot cache first (most likely to be used)
        if let Some(module) = self.hot_cache.read().get(id) {
            self.total_cache_hits.fetch_add(1, Ordering::Relaxed);
            self.record_execution(id);
            return Ok(module.clone());
        }

        // Check warm cache
        if let Some(module) = self.warm_cache.lock().get(id).cloned() {
            self.total_cache_hits.fetch_add(1, Ordering::Relaxed);
            self.record_execution(id);
            self.maybe_promote(id)?;
            return Ok(module);
        }

        // Check cold cache
        if let Some(module) = self.cold_cache.lock().get(id).cloned() {
            self.total_cache_hits.fetch_add(1, Ordering::Relaxed);
            self.record_execution(id);
            self.maybe_promote(id)?;
            return Ok(module);
        }

        // Cache miss - compile at cold tier
        self.total_cache_misses.fetch_add(1, Ordering::Relaxed);
        let module = self.compile_at_tier(id, bytecode, CompilationTier::Cold)?;

        // Store bytecode for later recompilation
        self.bytecode_store.write().insert(*id, bytecode.to_vec());

        // Cache at cold tier
        self.cold_cache.lock().put(*id, module.clone());

        // Initialize stats
        self.stats.write().insert(*id, ContractStats {
            execution_count: 1,
            tier: CompilationTier::Cold,
            last_executed: Instant::now(),
            ..Default::default()
        });

        Ok(module)
    }

    /// Compiles bytecode at a specific tier.
    pub fn compile_at_tier(
        &self,
        id: &ContractId,
        bytecode: &[u8],
        tier: CompilationTier,
    ) -> Result<Arc<TieredModule>, VmError> {
        if bytecode.len() > MAX_CONTRACT_SIZE {
            return Err(VmError::BytecodeTooLarge {
                size: bytecode.len(),
                max: MAX_CONTRACT_SIZE,
            });
        }

        let engine = self.engines.get(tier);
        let start = Instant::now();

        let module = Module::new(engine, bytecode)
            .map_err(|e| VmError::InvalidBytecode(e.to_string()))?;

        let compile_time = start.elapsed();
        self.total_compilations.fetch_add(1, Ordering::Relaxed);

        Ok(Arc::new(TieredModule::new(
            *id,
            bytecode.to_vec(),
            module,
            tier,
            compile_time,
        )))
    }

    /// Records an execution for tier promotion.
    fn record_execution(&self, id: &ContractId) {
        let mut stats = self.stats.write();
        if let Some(s) = stats.get_mut(id) {
            s.execution_count += 1;
            s.last_executed = Instant::now();
        }
    }

    /// Checks if a contract should be promoted to a higher tier.
    fn maybe_promote(&self, id: &ContractId) -> Result<(), VmError> {
        let stats = self.stats.read().get(id).cloned();

        if let Some(stats) = stats {
            let should_promote = match stats.tier {
                CompilationTier::Cold => stats.execution_count >= self.config.warm_threshold,
                CompilationTier::Warm => stats.execution_count >= self.config.hot_threshold,
                CompilationTier::Hot => false,
            };

            if should_promote {
                self.promote(id)?;
            }
        }

        Ok(())
    }

    /// Promotes a contract to the next tier.
    fn promote(&self, id: &ContractId) -> Result<(), VmError> {
        let current_tier = self.stats.read().get(id).map(|s| s.tier);

        if let Some(tier) = current_tier {
            let next_tier = match tier {
                CompilationTier::Cold => CompilationTier::Warm,
                CompilationTier::Warm => CompilationTier::Hot,
                CompilationTier::Hot => return Ok(()), // Already at max
            };

            // Get bytecode
            let bytecode = self.bytecode_store.read().get(id).cloned();

            if let Some(bytecode) = bytecode {
                let module = self.compile_at_tier(id, &bytecode, next_tier)?;

                // Move to appropriate cache
                match next_tier {
                    CompilationTier::Warm => {
                        self.cold_cache.lock().pop(id);
                        self.warm_cache.lock().put(*id, module);
                    }
                    CompilationTier::Hot => {
                        self.warm_cache.lock().pop(id);
                        self.hot_cache.write().insert(*id, module);
                    }
                    _ => {}
                }

                // Update stats
                if let Some(s) = self.stats.write().get_mut(id) {
                    s.tier = next_tier;
                }

                tracing::debug!(
                    "Promoted contract {} from {:?} to {:?}",
                    id,
                    tier,
                    next_tier
                );
            }
        }

        Ok(())
    }

    /// Prewarms the cache with frequently used contracts.
    pub fn prewarm(&self, contracts: Vec<(ContractId, Vec<u8>)>) -> Result<usize, VmError> {
        let mut count = 0;
        for (id, bytecode) in contracts {
            // Compile directly at warm tier for prewarming
            let module = self.compile_at_tier(&id, &bytecode, CompilationTier::Warm)?;
            self.warm_cache.lock().put(id, module);
            self.bytecode_store.write().insert(id, bytecode);
            self.stats.write().insert(id, ContractStats {
                tier: CompilationTier::Warm,
                ..Default::default()
            });
            count += 1;
        }
        Ok(count)
    }

    /// Evicts cold entries that haven't been used recently.
    pub fn evict_cold(&self, max_age: Duration) {
        let now = Instant::now();
        let mut to_evict = Vec::new();

        for (id, stats) in self.stats.read().iter() {
            if stats.tier == CompilationTier::Cold
                && now.duration_since(stats.last_executed) > max_age
            {
                to_evict.push(*id);
            }
        }

        for id in to_evict {
            self.cold_cache.lock().pop(&id);
            self.bytecode_store.write().remove(&id);
            self.stats.write().remove(&id);
        }
    }

    /// Returns compiler statistics.
    pub fn statistics(&self) -> CompilerStatistics {
        let stats = self.stats.read();

        let (cold_count, warm_count, hot_count) = stats.values().fold(
            (0, 0, 0),
            |(c, w, h), s| match s.tier {
                CompilationTier::Cold => (c + 1, w, h),
                CompilationTier::Warm => (c, w + 1, h),
                CompilationTier::Hot => (c, w, h + 1),
            },
        );

        CompilerStatistics {
            cold_contracts: cold_count,
            warm_contracts: warm_count,
            hot_contracts: hot_count,
            total_compilations: self.total_compilations.load(Ordering::Relaxed),
            cache_hits: self.total_cache_hits.load(Ordering::Relaxed),
            cache_misses: self.total_cache_misses.load(Ordering::Relaxed),
        }
    }

    /// Gets the engine for a specific tier (for external use).
    pub fn engine(&self, tier: CompilationTier) -> &Engine {
        self.engines.get(tier)
    }

    /// Gets the hot engine (for primary use with VmEngine).
    pub fn hot_engine(&self) -> &Engine {
        self.engines.get(CompilationTier::Hot)
    }
}

impl Default for TieredCompiler {
    fn default() -> Self {
        Self::new().expect("Failed to create tiered compiler")
    }
}

/// Statistics for the tiered compiler.
#[derive(Clone, Debug, Default)]
pub struct CompilerStatistics {
    /// Number of contracts at cold tier.
    pub cold_contracts: usize,
    /// Number of contracts at warm tier.
    pub warm_contracts: usize,
    /// Number of contracts at hot tier.
    pub hot_contracts: usize,
    /// Total compilations performed.
    pub total_compilations: u64,
    /// Cache hits.
    pub cache_hits: u64,
    /// Cache misses.
    pub cache_misses: u64,
}

impl CompilerStatistics {
    /// Returns cache hit ratio.
    pub fn hit_ratio(&self) -> f64 {
        let total = self.cache_hits + self.cache_misses;
        if total == 0 {
            0.0
        } else {
            self.cache_hits as f64 / total as f64
        }
    }
}

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

    fn make_contract(id: u8) -> ContractId {
        ContractId::from_bytes([id; 32])
    }

    // Minimal valid WASM module (empty module)
    fn minimal_wasm() -> Vec<u8> {
        // WAT: (module)
        vec![
            0x00, 0x61, 0x73, 0x6D, // WASM magic
            0x01, 0x00, 0x00, 0x00, // Version 1
        ]
    }

    #[test]
    fn test_tier_levels() {
        assert!(CompilationTier::Cold < CompilationTier::Warm);
        assert!(CompilationTier::Warm < CompilationTier::Hot);
    }

    #[test]
    fn test_tier_thresholds() {
        assert_eq!(CompilationTier::Cold.threshold(), 0);
        assert!(CompilationTier::Warm.threshold() > 0);
        assert!(CompilationTier::Hot.threshold() > CompilationTier::Warm.threshold());
    }

    #[test]
    fn test_compiler_creation() {
        let compiler = TieredCompiler::new();
        assert!(compiler.is_ok());
    }

    #[test]
    fn test_compile_minimal_wasm() {
        let compiler = TieredCompiler::new().unwrap();
        let id = make_contract(1);
        let wasm = minimal_wasm();

        let result = compiler.compile_at_tier(&id, &wasm, CompilationTier::Cold);
        assert!(result.is_ok());

        let module = result.unwrap();
        assert_eq!(module.tier, CompilationTier::Cold);
    }

    #[test]
    fn test_cache_hit() {
        let compiler = TieredCompiler::new().unwrap();
        let id = make_contract(1);
        let wasm = minimal_wasm();

        // First call - cache miss
        let _ = compiler.get_or_compile(&id, &wasm).unwrap();

        // Second call - cache hit
        let _ = compiler.get_or_compile(&id, &wasm).unwrap();

        let stats = compiler.statistics();
        assert_eq!(stats.cache_hits, 1);
        assert_eq!(stats.cache_misses, 1);
    }

    #[test]
    fn test_statistics() {
        let compiler = TieredCompiler::new().unwrap();
        let stats = compiler.statistics();

        assert_eq!(stats.cold_contracts, 0);
        assert_eq!(stats.warm_contracts, 0);
        assert_eq!(stats.hot_contracts, 0);
        assert_eq!(stats.hit_ratio(), 0.0);
    }

    #[test]
    fn test_prewarm() {
        let compiler = TieredCompiler::new().unwrap();
        let contracts = vec![
            (make_contract(1), minimal_wasm()),
            (make_contract(2), minimal_wasm()),
        ];

        let count = compiler.prewarm(contracts).unwrap();
        assert_eq!(count, 2);

        let stats = compiler.statistics();
        assert_eq!(stats.warm_contracts, 2);
    }
}
