from dataclasses import dataclass
from typing import Dict, List, Optional
import time
import asyncio
import aiohttp
import statistics
from concurrent.futures import ThreadPoolExecutor
import subprocess
from .utils import get_memory_usage, MemorySnapshot
import psutil
import gpustat

@dataclass
class ModelSharingMetrics:
    gpu_memory_sharing: Dict[str, float]  # Memory shared between models
    layer_allocation: Dict[str, List[int]]  # Layer distribution per model
    interference_score: float  # 0-1 score of inter-model interference
    memory_efficiency: float  # Ratio of shared vs total memory
    context_switches: int  # Number of model switches

@dataclass
class ConcurrencyMetrics:
    max_concurrent_requests: int
    average_latency: float
    latency_percentiles: Dict[str, float]  # p50, p90, p95, p99
    throughput: float  # requests per second
    queue_stats: Dict[str, float]  # queue length, wait time
    error_rate: float
    resource_contention: Dict[str, float]  # CPU, Memory, GPU contentions

@dataclass
class AdvancedMetrics:
    model_sharing: ModelSharingMetrics
    concurrency: ConcurrencyMetrics

class AdvancedBenchmark:
    def __init__(
        self,
        primary_model: str,
        secondary_models: List[str],
        num_concurrent_users: int = 10,
        test_duration: int = 60
    ):
        self.primary_model = primary_model
        self.secondary_models = secondary_models
        self.num_users = num_concurrent_users
        self.duration = test_duration
        self.base_prompt = "Explain the concept of quantum computing in simple terms."

    async def _send_concurrent_requests(
        self,
        model: str,
        num_requests: int
    ) -> List[Dict[str, float]]:
        """Send concurrent requests to the model and measure performance."""
        async with aiohttp.ClientSession() as session:
            async def single_request():
                start_time = time.time()
                try:
                    async with session.post(
                        "http://localhost:39281/v1/chat/completions",
                        json={
                            "model": model,
                            "messages": [{"role": "user", "content": self.base_prompt}]
                        }
                    ) as response:
                        await response.json()
                        return {
                            "latency": time.time() - start_time,
                            "success": response.status == 200
                        }
                except Exception as e:
                    return {
                        "latency": time.time() - start_time,
                        "success": False,
                        "error": str(e)
                    }

            tasks = [single_request() for _ in range(num_requests)]
            return await asyncio.gather(*tasks)

    def _measure_model_sharing(self) -> ModelSharingMetrics:
        """Measure metrics related to multiple models running together."""
        try:
            # Start primary model
            subprocess.run(["cortex", "run", "-d", self.primary_model], check=True)
            initial_memory = get_memory_usage()

            gpu_memory_sharing = {}
            layer_allocation = {self.primary_model: []}

            # Start secondary models and measure memory impact
            for model in self.secondary_models:
                subprocess.run(["cortex", "run", "-d", model], check=True)
                time.sleep(5)  # Wait for model to load

                current_memory = get_memory_usage()
                memory_delta = current_memory.gpu_memory_used - initial_memory.gpu_memory_used
                gpu_memory_sharing[model] = memory_delta
                layer_allocation[model] = []  # Would need cortex-specific API for layer info

            # Measure interference by running concurrent inference
            interference_scores = []
            for _ in range(5):  # Run multiple tests
                base_latency = asyncio.run(self._send_concurrent_requests(
                    self.primary_model, 1
                ))[0]["latency"]

                shared_latency = asyncio.run(self._send_concurrent_requests(
                    self.primary_model, 1
                ))[0]["latency"]

                interference_scores.append(shared_latency / base_latency)

            interference_score = statistics.mean(interference_scores)

            # Calculate memory efficiency
            total_memory = sum(gpu_memory_sharing.values())
            theoretical_memory = total_memory * len(self.secondary_models)
            memory_efficiency = total_memory / theoretical_memory if theoretical_memory > 0 else 1.0

            return ModelSharingMetrics(
                gpu_memory_sharing=gpu_memory_sharing,
                layer_allocation=layer_allocation,
                interference_score=interference_score,
                memory_efficiency=memory_efficiency,
                context_switches=len(self.secondary_models) * 2  # Approximate
            )
        finally:
            # Cleanup
            for model in self.secondary_models:
                subprocess.run(["cortex", "models", "stop", model], check=False)

    def _measure_concurrency(self) -> ConcurrencyMetrics:
        """Measure concurrent processing performance."""
        # Ensure primary model is running
        subprocess.run(["cortex", "run", "-d", self.primary_model], check=True)
        time.sleep(5)  # Wait for model to be ready

        try:
            # Run concurrent requests
            all_results = asyncio.run(self._send_concurrent_requests(
                self.primary_model,
                self.num_users
            ))

            # Calculate metrics
            latencies = [r["latency"] for r in all_results if r["success"]]
            success_count = sum(1 for r in all_results if r["success"])

            if not latencies:
                raise RuntimeError("No successful requests completed")

            # Calculate percentiles
            sorted_latencies = sorted(latencies)
            percentiles = {
                "p50": statistics.median(sorted_latencies),
                "p90": sorted_latencies[int(0.9 * len(sorted_latencies))],
                "p95": sorted_latencies[int(0.95 * len(sorted_latencies))],
                "p99": sorted_latencies[int(0.99 * len(sorted_latencies))]
            }

            return ConcurrencyMetrics(
                max_concurrent_requests=self.num_users,
                average_latency=statistics.mean(latencies),
                latency_percentiles=percentiles,
                throughput=success_count / max(latencies),
                queue_stats={
                    "avg_queue_length": self.num_users / 2,  # Approximation
                    "avg_wait_time": statistics.mean(latencies) / 2
                },
                error_rate=(len(all_results) - success_count) / len(all_results),
                resource_contention={
                    "cpu": psutil.cpu_percent(),
                    "memory": psutil.virtual_memory().percent,
                    "gpu": gpustat.GPUStatCollection.new_query()[0].utilization
                    if gpustat.GPUStatCollection.new_query() else 0
                }
            )
        finally:
            # Cleanup
            subprocess.run(["cortex", "models", "stop", self.primary_model], check=False)

    def benchmark_advanced(self) -> AdvancedMetrics:
        """Run advanced processing benchmarks."""
        model_sharing_metrics = self._measure_model_sharing()
        concurrency_metrics = self._measure_concurrency()

        return AdvancedMetrics(
            model_sharing=model_sharing_metrics,
            concurrency=concurrency_metrics
        )
