//! Advanced speculative decoder supporting multiple strategies

use super::config::SpeculativeConfig;
use super::lookahead::{LookaheadConfig, NgramPool};
use super::medusa::MedusaConfig;
use super::metrics::PerformanceMetrics;
use super::self_spec::SelfSpeculativeConfig;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};

/// Extended metrics for advanced speculative decoding
#[derive(Debug, Clone, Default)]
pub struct AdvancedMetrics {
    /// Base speculative metrics
    pub base_metrics: PerformanceMetrics,
    /// Tree attention speedup
    pub tree_attention_speedup: f32,
    /// Average accepted path length in tree
    pub avg_accepted_path_length: f32,
    /// Medusa head acceptance rates
    pub medusa_head_acceptance: Vec<f32>,
    /// Self-speculative early exit rates per layer
    pub early_exit_rates: Vec<f32>,
    /// N-gram cache hit rate
    pub ngram_hit_rate: f32,
}

/// Combined speculative decoding engine supporting multiple strategies
pub struct AdvancedSpeculativeDecoder {
    /// Basic speculative config
    pub base_config: SpeculativeConfig,
    /// Medusa config (if enabled)
    pub medusa_config: Option<MedusaConfig>,
    /// Self-speculative config (if enabled)
    pub self_spec_config: Option<SelfSpeculativeConfig>,
    /// Lookahead config (if enabled)
    pub lookahead_config: Option<LookaheadConfig>,
    /// N-gram pool for lookahead decoding
    pub ngram_pool: Option<Arc<RwLock<NgramPool>>>,
    /// Performance metrics
    pub metrics: Arc<Mutex<AdvancedMetrics>>,
}

impl AdvancedSpeculativeDecoder {
    /// Create a new advanced speculative decoder with basic config
    #[must_use]
    pub fn new(base_config: SpeculativeConfig) -> Self {
        Self {
            base_config,
            medusa_config: None,
            self_spec_config: None,
            lookahead_config: None,
            ngram_pool: None,
            metrics: Arc::new(Mutex::new(AdvancedMetrics::default())),
        }
    }

    /// Enable Medusa-style prediction
    #[must_use]
    pub fn with_medusa(mut self, config: MedusaConfig) -> Self {
        self.medusa_config = Some(config);
        self
    }

    /// Enable self-speculative decoding
    #[must_use]
    pub fn with_self_speculative(mut self, config: SelfSpeculativeConfig) -> Self {
        self.self_spec_config = Some(config);
        self
    }

    /// Enable lookahead decoding
    #[must_use]
    pub fn with_lookahead(mut self, config: LookaheadConfig) -> Self {
        let pool = NgramPool::new(config.ngram_size, config.max_pool_size);
        self.ngram_pool = Some(Arc::new(RwLock::new(pool)));
        self.lookahead_config = Some(config);
        self
    }

    /// Get current metrics
    pub async fn get_metrics(&self) -> AdvancedMetrics {
        self.metrics.lock().await.clone()
    }

    /// Add tokens to n-gram pool for lookahead decoding
    pub async fn observe_tokens(&self, tokens: &[u32]) {
        if let Some(ref pool) = self.ngram_pool {
            let mut pool = pool.write().await;
            pool.add_observation(tokens);
        }
    }

    /// Get lookahead predictions from n-gram pool
    pub async fn get_lookahead_predictions(&self, context: &[u32]) -> Vec<u32> {
        if let Some(ref pool) = self.ngram_pool {
            let pool = pool.read().await;
            pool.predict(context)
        } else {
            Vec::new()
        }
    }
}