from dataclasses import dataclass
from typing import Optional, List, Dict
import time
import json
import subprocess
from .utils import measure_time, get_memory_usage, MemorySnapshot
import statistics

@dataclass
class RuntimeMetrics:
    # Core Inference Metrics
    first_token_latency: float
    tokens_per_second: float
    token_latency_variance: float
    streaming_throughput: float

    # Context Handling
    context_window_size: int
    kv_cache_usage: float
    memory_per_token: float
    context_fill_rate: float

    # Batching Performance
    batch_throughput: Dict[int, float]  # batch_size -> tokens/sec
    optimal_batch_size: int
    queue_processing_time: float

    # Memory Usage
    peak_memory: MemorySnapshot
    memory_timeline: List[MemorySnapshot]

    # Hardware Utilization
    cpu_utilization: float
    gpu_utilization: Optional[float]
    temperature: Optional[float]
    power_consumption: Optional[float]

class RuntimeBenchmark:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.base_prompt = """
        You are a helpful AI assistant. Please help me understand the following:

        Write a detailed explanation of how a car engine works, including all major components
        and their functions. Make it technical but understandable.
        """
        self.streaming_prompt = "Tell me a story about space exploration in exactly 500 words."
        self.batch_prompts = [
            "Explain quantum computing in simple terms.",
            "Write a recipe for chocolate cake.",
            "Describe the water cycle.",
            "Explain how photosynthesis works.",
            "Describe the process of cellular respiration."
        ]

    def _measure_first_token_latency(self) -> float:
        start_time = time.perf_counter()
        # Run cortex command to get first token
        cmd = ["cortex", "run", "-d", self.model_name, "--first-token-only"]
        result = subprocess.run(cmd, capture_output=True, text=True)
        end_time = time.perf_counter()
        return end_time - start_time

    def _measure_streaming_performance(self) -> tuple[float, float]:
        tokens = []
        start_time = time.perf_counter()
        # Run cortex command with streaming
        cmd = ["cortex", "run", "-d", self.model_name, "--stream"]
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True
        )

        while True:
            output = process.stdout.readline()
            if output == '' and process.poll() is not None:
                break
            if output:
                tokens.append(output.strip())

        end_time = time.perf_counter()
        total_time = end_time - start_time
        total_tokens = len(tokens)

        return total_tokens / total_time, self._calculate_variance(tokens)

    def _measure_batch_performance(self, batch_sizes: List[int]) -> Dict[int, float]:
        results = {}
        for batch_size in batch_sizes:
            # Run cortex command with batch
            cmd = ["cortex", "run", "-d", self.model_name, f"--batch-size={batch_size}"]
            start_time = time.perf_counter()
            subprocess.run(cmd, capture_output=True, text=True)
            end_time = time.perf_counter()
            results[batch_size] = batch_size / (end_time - start_time)
        return results

    def _monitor_memory_usage(self, duration: int = 60) -> List[MemorySnapshot]:
        snapshots = []
        start_time = time.perf_counter()
        while time.perf_counter() - start_time < duration:
            snapshots.append(get_memory_usage())
            time.sleep(1)
        return snapshots

    def benchmark_runtime(self) -> RuntimeMetrics:
        # Measure first token latency
        first_token_latency = self._measure_first_token_latency()

        # Measure streaming performance
        tokens_per_second, latency_variance = self._measure_streaming_performance()

        # Measure batch performance
        batch_sizes = [1, 2, 4, 8, 16]
        batch_throughput = self._measure_batch_performance(batch_sizes)
        optimal_batch_size = max(batch_throughput.items(), key=lambda x: x[1])[0]

        # Monitor memory and hardware
        memory_timeline = self._monitor_memory_usage()
        peak_memory = max(memory_timeline, key=lambda x: x.ram_used)

        # Get hardware metrics
        hw_metrics = self._get_hardware_metrics()

        return RuntimeMetrics(
            first_token_latency=first_token_latency,
            tokens_per_second=tokens_per_second,
            token_latency_variance=latency_variance,
            streaming_throughput=tokens_per_second,  # Could be different metric
            context_window_size=self._get_context_window_size(),
            kv_cache_usage=self._get_kv_cache_usage(),
            memory_per_token=self._calculate_memory_per_token(),
            context_fill_rate=self._measure_context_fill_rate(),
            batch_throughput=batch_throughput,
            optimal_batch_size=optimal_batch_size,
            queue_processing_time=self._measure_queue_processing(),
            peak_memory=peak_memory,
            memory_timeline=memory_timeline,
            cpu_utilization=hw_metrics['cpu'],
            gpu_utilization=hw_metrics.get('gpu'),
            temperature=hw_metrics.get('temp'),
            power_consumption=hw_metrics.get('power')
        )

    # Helper methods to implement
    def _calculate_variance(self, tokens: List[str]) -> float:
        """Calculate variance in token generation times."""
        if not tokens:
            return 0.0
        token_times = [len(token) / 4 for token in tokens]  # Approximate time based on token length
        return statistics.variance(token_times) if len(token_times) > 1 else 0.0

    def _get_context_window_size(self) -> int:
        """Get context window size from model config."""
        try:
            result = subprocess.run(
                ["cortex", "models", "get", self.model_name, "--json"],
                capture_output=True,
                text=True,
                check=True
            )
            model_info = json.loads(result.stdout)
            return model_info.get("ctx_len", 2048)  # Default to 2048 if not found
        except Exception:
            return 2048

    def _get_kv_cache_usage(self) -> float:
        """Estimate KV cache usage."""
        memory_before = get_memory_usage()
        # Generate a long sequence to fill cache
        subprocess.run(
            ["cortex", "run", "-d", self.model_name],
            input="Generate a very long story about...",
            capture_output=True,
            text=True
        )
        memory_after = get_memory_usage()

        if memory_after.gpu_memory_used and memory_before.gpu_memory_used:
            return (memory_after.gpu_memory_used - memory_before.gpu_memory_used) / memory_after.gpu_memory_used * 100
        return 0.0

    def _calculate_memory_per_token(self) -> float:
        """Calculate average memory usage per token."""
        initial_memory = get_memory_usage()
        # Generate some text and count tokens
        result = subprocess.run(
            ["cortex", "run", "-d", self.model_name, "--json"],
            input=self.base_prompt,
            capture_output=True,
            text=True
        )
        final_memory = get_memory_usage()

        try:
            response = json.loads(result.stdout)
            num_tokens = response.get("usage", {}).get("total_tokens", 1)
            memory_delta = final_memory.ram_used - initial_memory.ram_used
            return memory_delta / num_tokens
        except Exception:
            return 0.0

    def _measure_context_fill_rate(self) -> float:
        """Measure how fast the context window fills."""
        start_time = time.time()
        result = subprocess.run(
            ["cortex", "run", "-d", self.model_name, "--json"],
            input=self.base_prompt * 10,  # Create a long input
            capture_output=True,
            text=True
        )
        end_time = time.time()

        try:
            response = json.loads(result.stdout)
            num_tokens = response.get("usage", {}).get("total_tokens", 0)
            return num_tokens / (end_time - start_time)
        except Exception:
            return 0.0

    def _measure_queue_processing(self) -> float:
        """Measure queue processing time."""
        start_time = time.time()
        # Submit multiple requests in quick succession
        for prompt in self.batch_prompts[:3]:
            subprocess.run(
                ["cortex", "run", "-d", self.model_name],
                input=prompt,
                capture_output=True
            )
        end_time = time.time()
        return (end_time - start_time) / 3

    def _get_hardware_metrics(self) -> Dict[str, float]:
        """Get hardware utilization metrics."""
        metrics = {
            'cpu': psutil.cpu_percent(),
            'gpu': None,
            'temp': None,
            'power': None
        }

        try:
            gpu_stats = gpustat.GPUStatCollection.new_query()
            if gpu_stats:
                metrics['gpu'] = gpu_stats[0].utilization
                metrics['temp'] = gpu_stats[0].temperature
                if hasattr(gpu_stats[0], 'power_draw'):
                    metrics['power'] = gpu_stats[0].power_draw
        except Exception:
            pass

        return metrics
