//! Configuration types for speculative decoding

use super::types::DraftModelType;
use crate::{InferenceError, InferenceResult};

/// Adaptive configuration for draft length optimization
#[derive(Debug, Clone)]
pub struct AdaptiveConfig {
    /// Minimum draft length to try
    pub min_length: usize,
    /// Maximum draft length to try
    pub max_length: usize,
    /// Target acceptance rate for adaptation
    pub target_acceptance_rate: f32,
    /// How aggressively to adjust draft length
    pub adjustment_factor: f32,
}

impl Default for AdaptiveConfig {
    fn default() -> Self {
        Self {
            min_length: 1,
            max_length: 8,
            target_acceptance_rate: 0.8,
            adjustment_factor: 0.1,
        }
    }
}

/// Configuration for speculative decoding
#[derive(Debug, Clone)]
pub struct SpeculativeConfig {
    /// Type of draft model to use
    pub draft_model_type: DraftModelType,
    /// Maximum number of tokens to generate in draft
    pub max_draft_tokens: usize,
    /// Minimum confidence threshold for acceptance
    pub acceptance_threshold: f32,
    /// Number of additional lookahead tokens
    pub lookahead_tokens: usize,
    /// Whether to track performance metrics
    pub performance_tracking: bool,
    /// Adaptive draft length configuration
    pub adaptive_config: Option<AdaptiveConfig>,
}

impl SpeculativeConfig {
    /// Create a new configuration with sensible defaults
    #[must_use]
    pub fn new() -> Self {
        Self {
            draft_model_type: DraftModelType::SmallModel("default".to_string()),
            max_draft_tokens: 4,
            acceptance_threshold: 0.7,
            lookahead_tokens: 1,
            performance_tracking: false,
            adaptive_config: None,
        }
    }

    /// Set the draft model type
    #[must_use]
    pub fn with_draft_model(mut self, model_type: DraftModelType) -> Self {
        self.draft_model_type = model_type;
        self
    }

    /// Set maximum draft tokens
    #[must_use]
    pub fn with_max_draft_tokens(mut self, max_tokens: usize) -> Self {
        self.max_draft_tokens = max_tokens;
        self
    }

    /// Set acceptance threshold
    #[must_use]
    pub fn with_acceptance_threshold(mut self, threshold: f32) -> Self {
        self.acceptance_threshold = threshold;
        self
    }

    /// Set lookahead tokens
    #[must_use]
    pub fn with_lookahead_tokens(mut self, lookahead: usize) -> Self {
        self.lookahead_tokens = lookahead;
        self
    }

    /// Enable performance tracking
    #[must_use]
    pub fn with_performance_tracking(mut self, enabled: bool) -> Self {
        self.performance_tracking = enabled;
        self
    }

    /// Enable adaptive draft length
    #[must_use]
    pub fn with_adaptive_draft_length(mut self, config: AdaptiveConfig) -> Self {
        self.adaptive_config = Some(config);
        self
    }

    /// Validate configuration parameters
    pub fn validate(&self) -> InferenceResult<()> {
        if self.max_draft_tokens == 0 {
            return Err(InferenceError::invalid_request(
                "max_draft_tokens must be greater than 0",
            ));
        }

        if !(0.0..=1.0).contains(&self.acceptance_threshold) {
            return Err(InferenceError::invalid_request(
                "acceptance_threshold must be between 0.0 and 1.0",
            ));
        }

        if let Some(ref adaptive) = self.adaptive_config {
            if adaptive.min_length == 0 || adaptive.min_length > adaptive.max_length {
                return Err(InferenceError::invalid_request(
                    "adaptive config: min_length must be > 0 and <= max_length",
                ));
            }
            if !(0.0..=1.0).contains(&adaptive.target_acceptance_rate) {
                return Err(InferenceError::invalid_request(
                    "adaptive config: target_acceptance_rate must be between 0.0 and 1.0",
                ));
            }
        }

        Ok(())
    }
}

impl Default for SpeculativeConfig {
    fn default() -> Self {
        Self::new()
    }
}