//! Configuration for llmux

use crate::policy::{CostAwarePolicy, FifoPolicy, SwitchPolicy, TimeSlicePolicy};
use anyhow::{Context, Result};
use onwards::target::Targets;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::time::Duration;

/// Top-level configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
    /// Models to manage
    pub models: HashMap<String, ModelConfig>,

    /// Switch policy configuration
    #[serde(default)]
    pub policy: PolicyConfig,

    /// Proxy port
    #[serde(default = "default_port")]
    pub port: u16,

    /// Metrics port (0 to disable)
    #[serde(default = "default_metrics_port")]
    pub metrics_port: u16,

    /// Command to use for spawning vLLM processes (default: "vllm")
    /// Can be overridden for testing with mock-vllm
    #[serde(default = "default_vllm_command")]
    pub vllm_command: String,

    /// CRIU/CUDA checkpoint configuration (enables sleep levels 3 and 4)
    #[serde(default)]
    pub checkpoint: Option<CheckpointConfig>,
}

/// Configuration for CUDA/CRIU-based checkpointing (sleep levels 3 and 4).
///
/// When present, enables `SleepLevel::Checkpoint` which uses `cuda-checkpoint`
/// and `criu` to snapshot the entire vLLM process to disk, freeing all GPU
/// memory while preserving full state (KV cache, CUDA graphs, warmed allocator).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointConfig {
    /// Path to the criu binary (default: "criu" on PATH)
    #[serde(default = "default_criu_path")]
    pub criu_path: String,

    /// Directory containing the CUDA checkpoint plugin (libcuda_plugin.so)
    #[serde(default = "default_cuda_plugin_dir")]
    pub cuda_plugin_dir: String,

    /// Base directory for checkpoint images (per-model subdirectories are created)
    #[serde(default = "default_images_dir")]
    pub images_dir: PathBuf,

    /// Path to cuda-checkpoint binary (default: "cuda-checkpoint" on PATH)
    #[serde(default = "default_cuda_checkpoint_path")]
    pub cuda_checkpoint_path: String,
}

fn default_vllm_command() -> String {
    "vllm".to_string()
}

fn default_criu_path() -> String {
    "criu".to_string()
}

fn default_cuda_plugin_dir() -> String {
    "/usr/lib/criu/".to_string()
}

fn default_images_dir() -> PathBuf {
    PathBuf::from("/tmp/llmux-checkpoints")
}

fn default_cuda_checkpoint_path() -> String {
    "cuda-checkpoint".to_string()
}

fn default_port() -> u16 {
    3000
}

fn default_metrics_port() -> u16 {
    9090
}

impl Config {
    /// Load configuration from a JSON file
    pub async fn from_file(path: &Path) -> Result<Self> {
        let contents = tokio::fs::read_to_string(path)
            .await
            .with_context(|| format!("Failed to read config file: {}", path.display()))?;

        serde_json::from_str(&contents)
            .with_context(|| format!("Failed to parse config file: {}", path.display()))
    }

    /// Validate configuration, warning about common misconfigurations.
    ///
    /// Checks that models using CudaSuspend (level 3) or Checkpoint (level 4)
    /// with tensor parallelism have the required vLLM flags.
    pub fn validate(&self) {
        use tracing::warn;

        for (name, model) in &self.models {
            let level = model.sleep_level;
            if level != 3 && level != 4 {
                continue;
            }

            // Parse --tensor-parallel-size from extra_args
            let tp = model.tensor_parallel_size();
            if tp <= 1 {
                continue;
            }

            let level_name = if level == 3 {
                "CudaSuspend"
            } else {
                "Checkpoint"
            };

            if !model.extra_args.contains(&"--enforce-eager".to_string()) {
                warn!(
                    model = %name,
                    "Model uses {} at TP={} but is missing --enforce-eager. \
                     CUDA graphs hold stale NCCL handles and will crash on resume.",
                    level_name, tp
                );
            }

            if !model
                .extra_args
                .contains(&"--disable-custom-all-reduce".to_string())
            {
                warn!(
                    model = %name,
                    "Model uses {} at TP={} but is missing --disable-custom-all-reduce. \
                     CustomAllReduce IPC buffers cannot survive cuda-checkpoint.",
                    level_name, tp
                );
            }
        }
    }

    /// Build onwards Targets from model configs
    pub fn build_onwards_targets(&self) -> Result<Targets> {
        use dashmap::DashMap;
        use onwards::load_balancer::ProviderPool;
        use onwards::target::Target;
        use std::sync::Arc;

        let targets_map: DashMap<String, ProviderPool> = DashMap::new();

        for (name, model_config) in &self.models {
            let url = format!("http://localhost:{}", model_config.port)
                .parse()
                .with_context(|| format!("Invalid URL for model {}", name))?;

            let target = Target {
                url,
                keys: None,
                onwards_key: None,
                onwards_model: Some(model_config.model_path.clone()),
                limiter: None,
                concurrency_limiter: None,
                upstream_auth_header_name: None,
                upstream_auth_header_prefix: None,
                response_headers: None,
                sanitize_response: false,
            };

            let pool = target.into_pool();
            targets_map.insert(name.clone(), pool);
        }

        Ok(Targets {
            targets: Arc::new(targets_map),
            key_rate_limiters: Arc::new(DashMap::new()),
            key_concurrency_limiters: Arc::new(DashMap::new()),
        })
    }
}

/// Configuration for a single model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Path to the model (HuggingFace model ID or local path)
    pub model_path: String,

    /// Port for this model's vLLM instance
    pub port: u16,

    /// Additional vLLM CLI arguments (e.g. `["--gpu-memory-utilization", "0.8"]`)
    #[serde(default)]
    pub extra_args: Vec<String>,

    /// Sleep level when hibernating:
    /// - 1: Offload weights to CPU RAM (fastest wake)
    /// - 2: Discard weights (slower wake, less RAM)
    /// - 3: CUDA suspend (process stays alive, GPU freed, state in host RAM)
    /// - 4: CRIU checkpoint (snapshot to disk, frees all GPU/CPU memory)
    /// - 5: Stop the vLLM process entirely (full restart on wake)
    #[serde(default = "default_sleep_level")]
    pub sleep_level: u8,
}

fn default_sleep_level() -> u8 {
    5
}

impl ModelConfig {
    /// Parse tensor parallel size from extra_args (default: 1)
    pub fn tensor_parallel_size(&self) -> usize {
        self.extra_args
            .windows(2)
            .find(|w| w[0] == "--tensor-parallel-size")
            .and_then(|w| w[1].parse().ok())
            .unwrap_or(1)
    }

    /// Build vLLM command line arguments
    pub fn vllm_args(&self) -> Vec<String> {
        let mut args = vec![
            "serve".to_string(),
            self.model_path.clone(),
            "--port".to_string(),
            self.port.to_string(),
            "--enable-sleep-mode".to_string(),
        ];

        args.extend(self.extra_args.clone());
        args
    }
}

/// Policy configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PolicyConfig {
    /// Policy type: "fifo" (default) or "cost_aware"
    #[serde(default = "default_policy_type")]
    pub policy_type: String,

    /// Request timeout in seconds
    #[serde(default = "default_request_timeout")]
    pub request_timeout_secs: u64,

    /// Whether to drain in-flight requests before switching
    #[serde(default = "default_drain_before_switch")]
    pub drain_before_switch: bool,

    /// Default sleep level (1-5)
    #[serde(default = "default_sleep_level")]
    pub sleep_level: u8,

    /// Minimum seconds a model must stay active before it can be put to sleep.
    /// Prevents rapid wake/sleep thrashing that can cause GPU page faults.
    #[serde(default = "default_min_active_secs")]
    pub min_active_secs: u64,

    /// (cost_aware) Coalescing window in milliseconds. On the first request
    /// for a non-active model, wait this long to collect more demand.
    #[serde(default = "default_coalesce_window_ms")]
    pub coalesce_window_ms: u64,

    /// (cost_aware) Amortization factor. Higher values require more pending
    /// requests before switching. 0.0 = always switch immediately.
    #[serde(default = "default_amortization_factor")]
    pub amortization_factor: f64,

    /// (cost_aware) Maximum wait in seconds before forcing a switch,
    /// regardless of cost calculations.
    #[serde(default = "default_max_wait_secs")]
    pub max_wait_secs: u64,

    /// (time_slice) Scheduler tick interval in milliseconds. Controls how
    /// often the background scheduler evaluates queue depths.
    #[serde(default = "default_tick_interval_ms")]
    pub tick_interval_ms: u64,

    /// (time_slice) Minimum quantum in milliseconds. Each model gets at
    /// least this much GPU time before the scheduler considers switching.
    #[serde(default = "default_min_quantum_ms")]
    pub min_quantum_ms: u64,
}

impl Default for PolicyConfig {
    fn default() -> Self {
        Self {
            policy_type: default_policy_type(),
            request_timeout_secs: default_request_timeout(),
            drain_before_switch: default_drain_before_switch(),
            sleep_level: default_sleep_level(),
            min_active_secs: default_min_active_secs(),
            coalesce_window_ms: default_coalesce_window_ms(),
            amortization_factor: default_amortization_factor(),
            max_wait_secs: default_max_wait_secs(),
            tick_interval_ms: default_tick_interval_ms(),
            min_quantum_ms: default_min_quantum_ms(),
        }
    }
}

fn default_policy_type() -> String {
    "fifo".to_string()
}

fn default_request_timeout() -> u64 {
    300
}

fn default_drain_before_switch() -> bool {
    true
}

fn default_min_active_secs() -> u64 {
    5
}

fn default_coalesce_window_ms() -> u64 {
    2000
}

fn default_amortization_factor() -> f64 {
    0.5
}

fn default_max_wait_secs() -> u64 {
    15
}

fn default_tick_interval_ms() -> u64 {
    500
}

fn default_min_quantum_ms() -> u64 {
    2000
}

impl PolicyConfig {
    /// Build a SwitchPolicy from this config.
    /// `model_names` is needed for cost_aware policy to pre-populate timing tables.
    pub fn build_policy(&self, model_names: &[String]) -> Box<dyn SwitchPolicy> {
        match self.policy_type.as_str() {
            "cost_aware" => Box::new(CostAwarePolicy::new(
                self.sleep_level,
                Duration::from_secs(self.request_timeout_secs),
                Duration::from_secs(self.min_active_secs),
                Duration::from_millis(self.coalesce_window_ms),
                self.amortization_factor,
                Duration::from_secs(self.max_wait_secs),
                model_names.to_vec(),
            )),
            "time_slice" => Box::new(TimeSlicePolicy::new(
                self.sleep_level,
                Duration::from_secs(self.request_timeout_secs),
                Duration::from_secs(self.min_active_secs),
                Duration::from_secs(self.max_wait_secs),
                Duration::from_millis(self.min_quantum_ms),
                Duration::from_millis(self.tick_interval_ms),
                model_names.to_vec(),
            )),
            _ => Box::new(FifoPolicy::new(
                self.sleep_level,
                Duration::from_secs(self.request_timeout_secs),
                self.drain_before_switch,
                Duration::from_secs(self.min_active_secs),
            )),
        }
    }
}

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

    #[test]
    fn test_parse_config() {
        let json = r#"{
            "models": {
                "llama": {
                    "model_path": "meta-llama/Llama-2-7b",
                    "port": 8001
                },
                "mistral": {
                    "model_path": "mistralai/Mistral-7B-v0.1",
                    "port": 8002,
                    "extra_args": ["--gpu-memory-utilization", "0.8"]
                }
            },
            "policy": {
                "request_timeout_secs": 30
            },
            "port": 3000
        }"#;

        let config: Config = serde_json::from_str(json).unwrap();
        assert_eq!(config.models.len(), 2);
        assert_eq!(config.models["llama"].port, 8001);
        assert_eq!(config.models["mistral"].extra_args.len(), 2);
        assert_eq!(config.policy.request_timeout_secs, 30);
    }

    #[test]
    fn test_vllm_args() {
        let config = ModelConfig {
            model_path: "meta-llama/Llama-2-7b".to_string(),
            port: 8001,
            extra_args: vec![
                "--tensor-parallel-size".to_string(),
                "2".to_string(),
                "--max-model-len".to_string(),
                "4096".to_string(),
            ],
            sleep_level: 1,
        };

        let args = config.vllm_args();
        assert!(args.contains(&"--enable-sleep-mode".to_string()));
        assert!(args.contains(&"--tensor-parallel-size".to_string()));
        assert!(args.contains(&"2".to_string()));
    }
}
