from dataclasses import dataclass
from typing import Dict, List, Optional
import time
import psutil
import gpustat
import subprocess
from .utils import get_memory_usage, MemorySnapshot

@dataclass
class ResourceMetrics:
    # Memory Management
    memory_usage_timeline: List[MemorySnapshot]
    peak_memory_usage: MemorySnapshot
    memory_growth_rate: float  # MB/s
    cache_hit_rate: Optional[float]  # percentage
    fragmentation_ratio: float  # ratio of used/allocated memory

    # Hardware Utilization
    cpu_usage_timeline: List[float]  # percentage per core
    cpu_frequency: float  # MHz
    cpu_temperature: Optional[float]  # Celsius

    gpu_metrics: Optional[Dict[str, Dict[str, float]]]  # GPU ID -> metrics
    pcie_bandwidth_usage: Optional[float]  # MB/s
    power_consumption: Optional[float]  # Watts

    # System Impact
    system_load_average: List[float]  # 1, 5, 15 minute averages
    io_operations: Dict[str, int]  # read/write operations
    network_usage: Dict[str, float]  # bytes sent/received

class ResourceBenchmark:
    def __init__(self, model_name: str, duration: int = 60):
        self.model_name = model_name
        self.duration = duration
        self.sampling_interval = 1  # seconds

    def _monitor_memory(self) -> tuple[List[MemorySnapshot], float]:
        """Monitor memory usage over time and calculate growth rate."""
        snapshots = []
        start_time = time.time()

        while time.time() - start_time < self.duration:
            snapshot = get_memory_usage()
            snapshots.append(snapshot)
            time.sleep(self.sampling_interval)

        # Calculate memory growth rate
        if len(snapshots) > 1:
            memory_growth = (snapshots[-1].ram_used - snapshots[0].ram_used)
            growth_rate = memory_growth / self.duration
        else:
            growth_rate = 0

        return snapshots, growth_rate

    def _monitor_cpu(self) -> tuple[List[float], float, Optional[float]]:
        """Monitor CPU usage, frequency, and temperature."""
        cpu_usage = []
        start_time = time.time()

        while time.time() - start_time < self.duration:
            cpu_usage.append(psutil.cpu_percent(interval=1, percpu=True))
            time.sleep(self.sampling_interval)

        # Get CPU frequency
        freq = psutil.cpu_freq()
        current_freq = freq.current if freq else 0

        # Try to get CPU temperature (not available on all systems)
        try:
            temps = psutil.sensors_temperatures()
            cpu_temp = max(temp.current for temp in temps.get('coretemp', []))
        except:
            cpu_temp = None

        return cpu_usage, current_freq, cpu_temp

    def _monitor_gpu(self) -> Optional[Dict[str, Dict[str, float]]]:
        """Monitor GPU metrics if available."""
        try:
            gpu_stats = gpustat.GPUStatCollection.new_query()
            gpu_metrics = {}

            for gpu in gpu_stats:
                gpu_metrics[str(gpu.index)] = {
                    'utilization': gpu.utilization,
                    'memory_used': gpu.memory_used,
                    'temperature': gpu.temperature,
                    'power_draw': gpu.power_draw if hasattr(gpu, 'power_draw') else None
                }

            return gpu_metrics
        except:
            return None

    def _measure_pcie_bandwidth(self) -> Optional[float]:
        """Measure PCIe bandwidth usage if possible."""
        try:
            # This is system dependent and might need adjustment
            cmd = ["nvidia-smi", "--query-gpu=pcie.rx_throughput,pcie.tx_throughput", "--format=csv,noheader,nounits"]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode == 0:
                rx, tx = map(float, result.stdout.strip().split(','))
                return rx + tx
            return None
        except:
            return None

    def _get_system_metrics(self) -> tuple[List[float], Dict[str, int], Dict[str, float]]:
        """Get system-wide metrics."""
        # Load averages
        load_avg = psutil.getloadavg()

        # IO operations
        io = psutil.disk_io_counters()
        io_ops = {
            'read_count': io.read_count if io else 0,
            'write_count': io.write_count if io else 0
        }

        # Network usage
        net = psutil.net_io_counters()
        net_usage = {
            'bytes_sent': net.bytes_sent if net else 0,
            'bytes_recv': net.bytes_recv if net else 0
        }

        return load_avg, io_ops, net_usage

    def benchmark_resources(self) -> ResourceMetrics:
        """Run resource utilization benchmarks."""
        # Start the model if not already running
        self._ensure_model_running()

        # Monitor memory
        memory_timeline, memory_growth = self._monitor_memory()
        peak_memory = max(memory_timeline, key=lambda x: x.ram_used)

        # Monitor CPU
        cpu_timeline, cpu_freq, cpu_temp = self._monitor_cpu()

        # Monitor GPU
        gpu_metrics = self._monitor_gpu()
        pcie_bandwidth = self._measure_pcie_bandwidth()

        # Get system metrics
        load_avg, io_ops, net_usage = self._get_system_metrics()

        # Calculate memory fragmentation (simplified)
        total_memory = psutil.virtual_memory().total / (1024 * 1024)  # MB
        used_memory = peak_memory.ram_used
        fragmentation_ratio = used_memory / total_memory

        return ResourceMetrics(
            memory_usage_timeline=memory_timeline,
            peak_memory_usage=peak_memory,
            memory_growth_rate=memory_growth,
            cache_hit_rate=None,  # Would need application-specific implementation
            fragmentation_ratio=fragmentation_ratio,

            cpu_usage_timeline=cpu_timeline,
            cpu_frequency=cpu_freq,
            cpu_temperature=cpu_temp,

            gpu_metrics=gpu_metrics,
            pcie_bandwidth_usage=pcie_bandwidth,
            power_consumption=None,  # Would need platform-specific implementation

            system_load_average=list(load_avg),
            io_operations=io_ops,
            network_usage=net_usage
        )

    def _ensure_model_running(self):
        """Make sure the model is running before benchmarking."""
        try:
            result = subprocess.run(
                ["cortex", "ps"],
                capture_output=True,
                text=True,
                check=True
            )
            if self.model_name not in result.stdout:
                subprocess.run(
                    ["cortex", "run", "-d", self.model_name],
                    check=True
                )
                time.sleep(10)  # Wait for model to load
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Failed to ensure model is running: {e}")
