from dataclasses import dataclass
from typing import Dict, List, Optional
import time
import asyncio
import aiohttp
import statistics
import psutil
import json
import websockets
from .utils import get_memory_usage, MemorySnapshot, BenchmarkExample

@dataclass
class APIMetrics:
    # Latency Metrics
    average_latency: float  # seconds
    latency_percentiles: Dict[str, float]  # p50, p90, p95, p99
    cold_call_latency: float  # first request latency
    warm_call_latency: float  # subsequent requests

    # Bandwidth Metrics
    average_request_size: float  # KB
    average_response_size: float  # KB
    bandwidth_utilization: float  # MB/s

    # Connection Metrics
    max_concurrent_connections: int
    connection_errors: int
    connection_timeouts: int
    average_connection_lifetime: float

    # WebSocket Performance
    websocket_latency: float
    streaming_stability: float  # 0-1 score
    message_throughput: float  # messages/second

    # Queue Metrics
    average_queue_length: float
    queue_wait_time: float
    queue_timeout_rate: float
    request_success_rate: float

class IntegrationBenchmark:
    def __init__(
        self,
        model_name: str,
        host: str = "localhost",
        port: int = 39281,
        num_concurrent: int = 10,
        test_duration: int = 60
    ):
        self.model_name = model_name
        self.base_url = f"http://{host}:{port}/v1"
        self.ws_url = f"ws://{host}:{port}/v1/stream"
        self.num_concurrent = num_concurrent
        self.duration = test_duration
        self.session_timeout = 30

    async def _measure_api_latency(self) -> Dict[str, float]:
        """Measure API latency with multiple requests."""
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.session_timeout)) as session:
            latencies = []
            errors = 0

            # First request (cold)
            cold_start = time.time()
            try:
                async with session.post(
                    f"{self.base_url}/chat/completions",
                    json={
                        "model": self.model_name,
                        "messages": [{"role": "user", "content": "Hello"}]
                    }
                ) as response:
                    await response.json()
                    cold_latency = time.time() - cold_start
            except Exception as e:
                print(f"Cold start error: {e}")
                cold_latency = float('inf')
                errors += 1

            # Warm requests
            for _ in range(self.num_concurrent):
                start = time.time()
                try:
                    async with session.post(
                        f"{self.base_url}/chat/completions",
                        json={
                            "model": self.model_name,
                            "messages": [{"role": "user", "content": "Quick test"}]
                        }
                    ) as response:
                        await response.json()
                        latencies.append(time.time() - start)
                except Exception:
                    errors += 1

            if not latencies:
                return {
                    "cold": cold_latency,
                    "warm": float('inf'),
                    "errors": errors
                }

            return {
                "cold": cold_latency,
                "warm": statistics.mean(latencies),
                "errors": errors
            }

    async def _measure_websocket_performance(self) -> Dict[str, float]:
        """Measure WebSocket streaming performance."""
        messages_received = 0
        start_time = time.time()
        latencies = []

        async with websockets.connect(self.ws_url) as websocket:
            # Send a streaming request
            await websocket.send(json.dumps({
                "model": self.model_name,
                "messages": [{"role": "user", "content": "Tell me a story"}],
                "stream": True
            }))

            while time.time() - start_time < 10:  # 10-second test
                try:
                    message_start = time.time()
                    message = await websocket.recv()
                    latencies.append(time.time() - message_start)
                    messages_received += 1
                except websockets.exceptions.ConnectionClosed:
                    break

        total_time = time.time() - start_time
        return {
            "latency": statistics.mean(latencies) if latencies else float('inf'),
            "throughput": messages_received / total_time if total_time > 0 else 0,
            "stability": 1 - (statistics.stdev(latencies) / statistics.mean(latencies)) if latencies else 0
        }

    async def _measure_queue_performance(self) -> Dict[str, float]:
        """Measure request queue performance."""
        queue_metrics = {
            "lengths": [],
            "wait_times": [],
            "timeouts": 0,
            "successes": 0
        }

        async def send_request(session, delay: float = 0):
            if delay:
                await asyncio.sleep(delay)

            start_time = time.time()
            try:
                async with session.post(
                    f"{self.base_url}/chat/completions",
                    json={
                        "model": self.model_name,
                        "messages": [{"role": "user", "content": "Test request"}]
                    }
                ) as response:
                    await response.json()
                    queue_metrics["wait_times"].append(time.time() - start_time)
                    queue_metrics["successes"] += 1
            except asyncio.TimeoutError:
                queue_metrics["timeouts"] += 1

        async with aiohttp.ClientSession() as session:
            # Send requests with varying delays
            tasks = [
                send_request(session, delay=i*0.5)
                for i in range(self.num_concurrent)
            ]

            for task in asyncio.as_completed(tasks):
                try:
                    await task
                except Exception as e:
                    print(f"Queue test error: {e}")

        return {
            "avg_wait": statistics.mean(queue_metrics["wait_times"]) if queue_metrics["wait_times"] else float('inf'),
            "timeout_rate": queue_metrics["timeouts"] / self.num_concurrent,
            "success_rate": queue_metrics["successes"] / self.num_concurrent
        }

    async def _measure_bandwidth(self) -> Dict[str, float]:
        """Measure bandwidth usage with varying request sizes."""
        sizes = []
        network_before = psutil.net_io_counters()

        async with aiohttp.ClientSession() as session:
            # Send requests with varying content sizes
            for length in [100, 500, 1000]:
                content = "Test " * length
                try:
                    async with session.post(
                        f"{self.base_url}/chat/completions",
                        json={
                            "model": self.model_name,
                            "messages": [{"role": "user", "content": content}]
                        }
                    ) as response:
                        response_data = await response.json()
                        sizes.append({
                            "request": len(content),
                            "response": len(str(response_data))
                        })
                except Exception as e:
                    print(f"Bandwidth test error: {e}")

        network_after = psutil.net_io_counters()
        total_bytes = (network_after.bytes_sent + network_after.bytes_recv -
                      network_before.bytes_sent - network_before.bytes_recv)

        return {
            "avg_request_size": statistics.mean([s["request"] for s in sizes]) / 1024 if sizes else 0,
            "avg_response_size": statistics.mean([s["response"] for s in sizes]) / 1024 if sizes else 0,
            "bandwidth": total_bytes / (1024 * 1024)  # MB
        }

    async def benchmark_integration(self) -> APIMetrics:
        """Run all integration benchmarks."""
        # Measure API latency
        latency_results = await self._measure_api_latency()

        # Measure WebSocket performance
        ws_results = await self._measure_websocket_performance()

        # Measure queue performance
        queue_results = await self._measure_queue_performance()

        # Measure bandwidth
        bandwidth_results = await self._measure_bandwidth()

        # Calculate percentiles from all latencies
        all_latencies = [
            result["latency"]
            for result in await asyncio.gather(*[
                self._measure_api_latency()
                for _ in range(5)
            ])
            if result["latency"] != float('inf')
        ]

        percentiles = {
            "p50": statistics.median(all_latencies),
            "p90": all_latencies[int(0.9 * len(all_latencies))],
            "p95": all_latencies[int(0.95 * len(all_latencies))],
            "p99": all_latencies[int(0.99 * len(all_latencies))]
        } if all_latencies else {
            "p50": float('inf'),
            "p90": float('inf'),
            "p95": float('inf'),
            "p99": float('inf')
        }

        return APIMetrics(
            average_latency=statistics.mean(all_latencies) if all_latencies else float('inf'),
            latency_percentiles=percentiles,
            cold_call_latency=latency_results["cold"],
            warm_call_latency=latency_results["warm"],

            average_request_size=bandwidth_results["avg_request_size"],
            average_response_size=bandwidth_results["avg_response_size"],
            bandwidth_utilization=bandwidth_results["bandwidth"],

            max_concurrent_connections=self.num_concurrent,
            connection_errors=latency_results["errors"],
            connection_timeouts=0,  # Updated during tests
            average_connection_lifetime=queue_results["avg_wait"],

            websocket_latency=ws_results["latency"],
            streaming_stability=ws_results["stability"],
            message_throughput=ws_results["throughput"],

            average_queue_length=self.num_concurrent / 2,  # Approximation
            queue_wait_time=queue_results["avg_wait"],
            queue_timeout_rate=queue_results["timeout_rate"],
            request_success_rate=queue_results["success_rate"]
        )
