from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import time
import psutil
import gpustat
import statistics
import subprocess
import json
from .utils import get_memory_usage, MemorySnapshot
from rich.console import Console

console = Console()

@dataclass
class StabilityMetrics:
    # Performance Degradation
    initial_throughput: float  # tokens/second
    final_throughput: float   # tokens/second
    degradation_rate: float   # percentage/hour
    response_time_trend: List[float]  # List of response times

    # Memory Management
    memory_leak_detected: bool
    memory_growth_rate: float  # MB/hour
    memory_pattern: List[MemorySnapshot]
    baseline_memory: float     # MB
    peak_memory: float        # MB

    # Error Analysis
    error_rate: float         # errors/hour
    error_types: Dict[str, int]  # error type -> count
    recovery_time: float      # average seconds to recover
    crash_count: int

    # Thermal Management
    temperature_profile: List[float]  # CPU/GPU temperatures
    thermal_throttling: bool
    throttling_impact: float  # performance impact percentage

    # System Stability
    uptime: float            # hours
    restart_count: int
    system_load: List[float]
    resource_exhaustion_events: int

class StabilityBenchmark:
    def __init__(
        self,
        model_name: str,
        duration_hours: float = 1.0,
        sampling_interval: int = 60,  # seconds
        stress_test: bool = True
    ):
        self.model_name = model_name
        self.duration = int(duration_hours * 3600)  # convert to seconds
        self.interval = sampling_interval
        self.stress_test = stress_test
        self.base_prompt = "Explain the concept of artificial intelligence."

    def _monitor_performance(self) -> Tuple[List[float], List[MemorySnapshot], List[Dict]]:
        """Monitor performance metrics over time."""
        performance_metrics = []
        memory_snapshots = []
        error_logs = []
        start_time = time.time()

        while time.time() - start_time < self.duration:
            try:
                # Measure response time and throughput
                perf_metric = self._measure_single_performance()
                performance_metrics.append(perf_metric)

                # Monitor memory
                memory_snapshots.append(get_memory_usage())

                # Check for errors
                error_log = self._check_system_health()
                if error_log:
                    error_logs.append(error_log)

                # Wait for next sampling interval
                time.sleep(self.interval)

            except Exception as e:
                error_logs.append({
                    "timestamp": time.time(),
                    "type": type(e).__name__,
                    "message": str(e)
                })

        return performance_metrics, memory_snapshots, error_logs

    def _measure_single_performance(self) -> float:
        """Measure performance for a single request."""
        try:
            start_time = time.time()
            result = subprocess.run(
                ["cortex", "run", "-d", self.model_name, "--json"],
                input=self.base_prompt,
                capture_output=True,
                text=True,
                check=True
            )
            response_time = time.time() - start_time

            # Parse response to get token count
            response = json.loads(result.stdout)
            tokens = response.get("usage", {}).get("completion_tokens", 0)

            return tokens / response_time  # tokens per second

        except Exception as e:
            print(f"Performance measurement error: {e}")
            return 0.0

    def _check_system_health(self) -> Optional[Dict]:
        """Check system health and record any issues."""
        try:
            cpu_percent = psutil.cpu_percent(interval=1)
            memory = psutil.virtual_memory()

            # Get GPU metrics if available
            gpu_temp = None
            try:
                gpu_stats = gpustat.GPUStatCollection.new_query()
                if gpu_stats:
                    gpu_temp = gpu_stats[0].temperature
            except Exception:
                pass

            # Check for potential issues
            issues = []
            if cpu_percent > 95:
                issues.append("CPU_OVERLOAD")
            if memory.percent > 90:
                issues.append("MEMORY_CRITICAL")
            if gpu_temp and gpu_temp > 80:
                issues.append("GPU_OVERHEATING")

            if issues:
                return {
                    "timestamp": time.time(),
                    "issues": issues,
                    "cpu": cpu_percent,
                    "memory": memory.percent,
                    "gpu_temp": gpu_temp
                }

            return None

        except Exception as e:
            return {
                "timestamp": time.time(),
                "type": "MONITORING_ERROR",
                "message": str(e)
            }

    def _detect_memory_leak(self, snapshots: List[MemorySnapshot]) -> Tuple[bool, float]:
        """Analyze memory pattern to detect leaks."""
        if len(snapshots) < 2:
            return False, 0.0

        # Calculate memory growth rate
        memory_values = [s.ram_used for s in snapshots]
        time_hours = self.duration / 3600

        # Linear regression would be more accurate, but this is a simple approximation
        growth_rate = (memory_values[-1] - memory_values[0]) / time_hours

        # Consider it a leak if memory grows more than 10% per hour consistently
        baseline_memory = memory_values[0]
        is_leak = growth_rate > (baseline_memory * 0.1)

        return is_leak, growth_rate

    def _analyze_thermal_behavior(self, logs: List[Dict]) -> Tuple[bool, float]:
        """Analyze thermal behavior and its impact on performance."""
        thermal_events = [
            log for log in logs
            if "GPU_OVERHEATING" in log.get("issues", [])
        ]

        if not thermal_events:
            return False, 0.0

        # Calculate performance impact during thermal events
        normal_perf = statistics.mean([
            log.get("performance", 0) for log in logs
            if "GPU_OVERHEATING" not in log.get("issues", [])
        ])

        thermal_perf = statistics.mean([
            log.get("performance", 0) for log in logs
            if "GPU_OVERHEATING" in log.get("issues", [])
        ])

        impact = (normal_perf - thermal_perf) / normal_perf if normal_perf > 0 else 0
        return bool(thermal_events), impact

    def benchmark_stability(self) -> StabilityMetrics:
        """Run comprehensive stability benchmarks."""
        console.print("[yellow]Starting stability benchmark...")
        start_time = time.time()

        # Monitor performance over time
        perf_metrics, memory_snapshots, error_logs = self._monitor_performance()

        # Analyze performance degradation
        initial_perf = statistics.mean(perf_metrics[:5]) if perf_metrics else 0
        final_perf = statistics.mean(perf_metrics[-5:]) if perf_metrics else 0
        degradation = (initial_perf - final_perf) / initial_perf if initial_perf > 0 else 0

        # Analyze memory patterns
        memory_leak, growth_rate = self._detect_memory_leak(memory_snapshots)

        # Analyze errors
        error_types = {}
        for log in error_logs:
            error_type = log.get("type", "UNKNOWN")
            error_types[error_type] = error_types.get(error_type, 0) + 1

        # Analyze thermal behavior
        thermal_throttling, throttle_impact = self._analyze_thermal_behavior(error_logs)

        # Calculate temperature profile
        temp_profile = []
        try:
            gpu_stats = gpustat.GPUStatCollection.new_query()
            if gpu_stats:
                temp_profile = [gpu.temperature for gpu in gpu_stats]
        except Exception:
            pass

        return StabilityMetrics(
            # Performance Degradation
            initial_throughput=initial_perf,
            final_throughput=final_perf,
            degradation_rate=degradation * 100,  # convert to percentage
            response_time_trend=perf_metrics,

            # Memory Management
            memory_leak_detected=memory_leak,
            memory_growth_rate=growth_rate,
            memory_pattern=memory_snapshots,
            baseline_memory=memory_snapshots[0].ram_used if memory_snapshots else 0,
            peak_memory=max(s.ram_used for s in memory_snapshots) if memory_snapshots else 0,

            # Error Analysis
            error_rate=len(error_logs) / (self.duration / 3600),
            error_types=error_types,
            recovery_time=0.0,  # Would need more sophisticated recovery tracking
            crash_count=sum(1 for log in error_logs if "CRASH" in log.get("type", "")),

            # Thermal Management
            temperature_profile=temp_profile,
            thermal_throttling=thermal_throttling,
            throttling_impact=throttle_impact * 100,  # convert to percentage

            # System Stability
            uptime=(time.time() - start_time) / 3600,  # convert to hours
            restart_count=sum(1 for log in error_logs if "RESTART" in log.get("type", "")),
            system_load=psutil.getloadavg(),
            resource_exhaustion_events=sum(
                1 for log in error_logs
                if any(issue in ["CPU_OVERLOAD", "MEMORY_CRITICAL", "GPU_OVERHEATING"]
                      for issue in log.get("issues", []))
            )
        )
