import time
import json
import subprocess
from pathlib import Path
from dataclasses import dataclass
from typing import Optional, Dict, Any
from .utils import measure_time, get_memory_usage, MemorySnapshot
import yaml

@dataclass
class InitializationMetrics:
    # Required fields (no defaults) first
    model_path: Path
    model_name: str
    model_size: int
    quantization: str
    initial_memory: MemorySnapshot
    peak_memory: MemorySnapshot
    disk_to_ram_time: float
    cold_start_time: float
    warm_start_time: float
    memory_spike: float

    # Optional fields (with defaults) last
    model_switch_time: Optional[float] = None
    gpu_load_time: Optional[float] = None
    multi_gpu_times: Optional[Dict[str, float]] = None
    gpu_memory_allocation: Optional[Dict[str, float]] = None

@dataclass
class ModelConfig:
    id: str
    model: str
    name: str
    version: int
    files: list[str]
    stop: list[str]
    size: int
    stream: bool
    top_p: float
    temperature: float
    frequency_penalty: float
    presence_penalty: float
    max_tokens: int
    seed: int
    dynatemp_range: float
    dynatemp_exponent: float
    top_k: int
    min_p: float
    tfs_z: float
    typ_p: float
    repeat_last_n: int
    repeat_penalty: float
    mirostat: bool
    mirostat_tau: float
    mirostat_eta: float
    penalize_nl: bool
    ignore_eos: bool
    n_probs: int
    min_keep: int
    engine: str
    prompt_template: str
    ctx_len: int
    n_parallel: int
    ngl: int

def load_model_config(yaml_path: str) -> ModelConfig:
    with open(yaml_path, 'r') as f:
        config = yaml.safe_load(f)
        return ModelConfig(**config)


@dataclass
class ModelLoadMetrics:
    model_path: Optional[Path]
    model_name: str
    disk_to_ram_time: float
    gpu_load_time: Optional[float]
    multi_gpu_load_time: Optional[float]
    peak_memory: MemorySnapshot
    initial_memory: MemorySnapshot
    warm_start_time: float
    cold_start_time: float

class ModelInitializationBenchmark:
    def __init__(self, model_name: str):
        if ":" not in model_name:
            raise ValueError("Model name must be in format 'model:quantization'")
        self.model_name = model_name
        try:
            self.model_id, self.quantization = model_name.split(":")
        except ValueError:
            raise ValueError("Invalid model name format. Expected 'model:quantization'")

    def _run_cortex_command(self, command: list[str]) -> dict:
        """Execute cortex command and parse output."""
        try:
            result = subprocess.run(
                command,
                capture_output=True,
                text=True,
                check=True
            )
            # Only try to parse as JSON if we expect JSON output
            if any(cmd in ["get", "list"] for cmd in command):
                return json.loads(result.stdout)
            return {"status": "success", "output": result.stdout}
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Cortex command failed: {e.stderr}")
        except json.JSONDecodeError:
            # Return raw output if not JSON
            return {"status": "success", "output": result.stdout}

    @measure_time
    def _cold_start(self) -> Dict[str, Any]:
        """Perform cold start by stopping and starting the model"""
        # Stop any running instances
        self._run_cortex_command(["cortex", "stop"])

        # Clear GPU memory
        self._run_cortex_command(["cortex", "hardware", "activate", "--gpus", "[]"])

        initial_memory = get_memory_usage()

        # Start model and capture metrics
        result = self._run_cortex_command([
            "cortex", "run", "-d",
            self.model_name,
            "--verbose"
        ])

        peak_memory = get_memory_usage()

        return {
            "initial_memory": initial_memory,
            "peak_memory": peak_memory,
            "result": result
        }

    @measure_time
    def _warm_start(self) -> float:
        """Measure warm start time by reloading already loaded model"""
        self._run_cortex_command(["cortex", "models", "stop", self.model_name])
        return self._run_cortex_command(["cortex", "models", "start", self.model_name])

    @measure_time
    def _measure_model_switch(self, other_model: str) -> float:
        """Measure time to switch between models"""
        self._run_cortex_command(["cortex", "models", "start", other_model])
        return self._run_cortex_command(["cortex", "models", "start", self.model_name])

    def benchmark_initialization(self) -> InitializationMetrics:
        # Cold start benchmarking
        cold_start_data, cold_time = self._cold_start()

        # Warm start benchmarking
        _, warm_time = self._warm_start()

        # Model switching (if another model is available)
        model_switch_time = None
        try:
            other_model = "tinyllama:1b-gguf-q4"
            _, switch_time = self._measure_model_switch(other_model)
            model_switch_time = switch_time
        except Exception:
            pass

        # Get model info
        model_info = self._run_cortex_command([
            "cortex", "models", "get", self.model_name
        ])

        memory_spike = cold_start_data["peak_memory"].ram_used - \
                        cold_start_data["initial_memory"].ram_used

        return InitializationMetrics(
            model_path=Path(model_info.get("path", "")),
            model_name=self.model_name,
            model_size=model_info.get("size", 0),
            quantization=self.quantization,
            initial_memory=cold_start_data["initial_memory"],
            peak_memory=cold_start_data["peak_memory"],
            disk_to_ram_time=cold_time - warm_time,
            cold_start_time=cold_time,
            warm_start_time=warm_time,
            memory_spike=memory_spike,
            model_switch_time=model_switch_time,
            gpu_load_time=None,  # Will be populated if GPU metrics are available
            multi_gpu_times=None,
            gpu_memory_allocation=None
        )
