//! Medusa-style prediction head configuration

use super::tree::CandidateTree;
use crate::InferenceResult;

/// Medusa-style prediction head configuration
/// Multiple heads predict tokens at different positions simultaneously
#[derive(Debug, Clone)]
pub struct MedusaConfig {
    /// Number of prediction heads (each predicts one future position)
    pub num_heads: usize,
    /// Top-k candidates per head
    pub top_k_per_head: usize,
    /// Temperature for sampling from heads
    pub temperature: f32,
    /// Whether to use tree attention for verification
    pub use_tree_attention: bool,
    /// Typical acceptance threshold for Medusa heads
    pub typical_acceptance: f32,
}

impl Default for MedusaConfig {
    fn default() -> Self {
        Self {
            num_heads: 4,
            top_k_per_head: 10,
            temperature: 1.0,
            use_tree_attention: true,
            typical_acceptance: 0.6,
        }
    }
}

impl MedusaConfig {
    /// Create config for fast inference (fewer heads, lower k)
    #[must_use]
    pub fn fast() -> Self {
        Self {
            num_heads: 2,
            top_k_per_head: 5,
            temperature: 0.8,
            use_tree_attention: true,
            typical_acceptance: 0.5,
        }
    }

    /// Create config for high quality (more heads, higher k)
    #[must_use]
    pub fn high_quality() -> Self {
        Self {
            num_heads: 5,
            top_k_per_head: 15,
            temperature: 1.0,
            use_tree_attention: true,
            typical_acceptance: 0.7,
        }
    }
}

/// Trait for Medusa-style multi-head draft prediction
#[async_trait::async_trait]
pub trait MedusaDraftModel: Send + Sync {
    /// Generate predictions from all heads simultaneously
    /// Returns Vec<Vec<(`token_id`, probability)>> - one vec per head
    async fn predict_all_heads(
        &self,
        context: &[u32],
        config: &MedusaConfig,
    ) -> InferenceResult<Vec<Vec<(u32, f32)>>>;

    /// Build a candidate tree from head predictions
    async fn build_candidate_tree(
        &self,
        context: &[u32],
        config: &MedusaConfig,
    ) -> InferenceResult<CandidateTree> {
        let head_predictions = self.predict_all_heads(context, config).await?;

        let mut tree = CandidateTree::new();

        if head_predictions.is_empty() {
            return Ok(tree);
        }

        // Add root candidates from first head
        let first_head = &head_predictions[0];
        let root_indices: Vec<usize> = first_head
            .iter()
            .take(config.top_k_per_head)
            .map(|(token_id, prob)| tree.add_root(*token_id, *prob))
            .collect();

        // For each subsequent head, add children to existing nodes
        let mut current_level = root_indices;

        for head_preds in head_predictions.iter().skip(1) {
            let mut next_level = Vec::new();

            for &parent_idx in &current_level {
                for (token_id, prob) in head_preds.iter().take(config.top_k_per_head) {
                    let child_idx = tree.add_child(parent_idx, *token_id, *prob);
                    next_level.push(child_idx);
                }
            }

            current_level = next_level;
        }

        tree.total_candidates = tree.get_all_paths().len();

        Ok(tree)
    }
}