"""
Latency Optimization Utilities

Provides:
- Connection pooling
- Response caching
- Async optimizations
- Health check endpoints
- Warm-up routines
"""
from __future__ import annotations

import asyncio
import os
import time
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any, Callable, Optional, TypeVar

from app.observability.logging import get_logger


_logger = get_logger("latency")

T = TypeVar("T")


@dataclass
class ConnectionPool:
    """Generic async connection pool."""
    
    factory: Callable
    min_size: int = 2
    max_size: int = 10
    timeout: float = 5.0
    _pool: asyncio.Queue = field(default_factory=lambda: asyncio.Queue())
    _size: int = 0
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    
    async def initialize(self) -> None:
        """Pre-warm the pool with minimum connections."""
        async with self._lock:
            for _ in range(self.min_size):
                if self._size < self.max_size:
                    conn = await self.factory()
                    await self._pool.put(conn)
                    self._size += 1
        _logger.info("connection_pool_initialized", size=self._size)
    
    @asynccontextmanager
    async def acquire(self):
        """Acquire a connection from the pool."""
        conn = None
        try:
            # Try to get from pool first
            try:
                conn = self._pool.get_nowait()
            except asyncio.QueueEmpty:
                # Create new if under limit
                async with self._lock:
                    if self._size < self.max_size:
                        conn = await asyncio.wait_for(
                            self.factory(),
                            timeout=self.timeout,
                        )
                        self._size += 1
                    else:
                        # Wait for one to be returned
                        conn = await asyncio.wait_for(
                            self._pool.get(),
                            timeout=self.timeout,
                        )
            yield conn
        finally:
            if conn is not None:
                try:
                    await self._pool.put(conn)
                except Exception:
                    async with self._lock:
                        self._size -= 1
    
    async def close(self) -> None:
        """Close all connections in the pool."""
        async with self._lock:
            while not self._pool.empty():
                try:
                    conn = self._pool.get_nowait()
                    if hasattr(conn, "close"):
                        await conn.close()
                except Exception:
                    pass
            self._size = 0


class ResponseCache:
    """
    Simple in-memory cache for API responses.
    
    Uses TTL-based expiration and LRU eviction.
    """
    
    def __init__(self, max_size: int = 1000, default_ttl: int = 60):
        self.max_size = max_size
        self.default_ttl = default_ttl
        self._cache: dict[str, tuple[Any, float]] = {}
        self._lock = asyncio.Lock()
    
    async def get(self, key: str) -> Optional[Any]:
        """Get value from cache if not expired."""
        async with self._lock:
            if key not in self._cache:
                return None
            value, expiry = self._cache[key]
            if time.time() > expiry:
                del self._cache[key]
                return None
            return value
    
    async def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
        """Set value in cache with TTL."""
        if ttl is None:
            ttl = self.default_ttl
        
        async with self._lock:
            # Evict oldest if at capacity
            if len(self._cache) >= self.max_size:
                oldest_key = min(self._cache, key=lambda k: self._cache[k][1])
                del self._cache[oldest_key]
            
            self._cache[key] = (value, time.time() + ttl)
    
    async def delete(self, key: str) -> None:
        """Delete value from cache."""
        async with self._lock:
            self._cache.pop(key, None)
    
    async def clear(self) -> None:
        """Clear all cached values."""
        async with self._lock:
            self._cache.clear()
    
    def stats(self) -> dict:
        """Get cache statistics."""
        return {
            "size": len(self._cache),
            "max_size": self.max_size,
        }


# Global cache instance
_response_cache = ResponseCache()


def cache_key(*args, **kwargs) -> str:
    """Generate cache key from arguments."""
    import hashlib
    import json
    
    key_data = json.dumps({"args": args, "kwargs": kwargs}, sort_keys=True, default=str)
    return hashlib.md5(key_data.encode()).hexdigest()


def cached(ttl: int = 60):
    """Decorator for caching async function results."""
    def decorator(func: Callable):
        async def wrapper(*args, **kwargs):
            key = f"{func.__name__}:{cache_key(*args, **kwargs)}"
            
            # Try cache first
            cached_value = await _response_cache.get(key)
            if cached_value is not None:
                return cached_value
            
            # Call function and cache result
            result = await func(*args, **kwargs)
            await _response_cache.set(key, result, ttl)
            return result
        
        return wrapper
    return decorator


@dataclass
class LatencyTracker:
    """Track and report latency metrics."""
    
    name: str
    _samples: list[float] = field(default_factory=list)
    _max_samples: int = 1000
    
    def record(self, duration_ms: float) -> None:
        """Record a latency sample."""
        self._samples.append(duration_ms)
        if len(self._samples) > self._max_samples:
            self._samples = self._samples[-self._max_samples:]
    
    @property
    def avg(self) -> float:
        """Average latency in ms."""
        return sum(self._samples) / len(self._samples) if self._samples else 0
    
    @property
    def p50(self) -> float:
        """50th percentile latency."""
        return self._percentile(50)
    
    @property
    def p95(self) -> float:
        """95th percentile latency."""
        return self._percentile(95)
    
    @property
    def p99(self) -> float:
        """99th percentile latency."""
        return self._percentile(99)
    
    def _percentile(self, p: int) -> float:
        if not self._samples:
            return 0
        sorted_samples = sorted(self._samples)
        idx = int(len(sorted_samples) * p / 100)
        return sorted_samples[min(idx, len(sorted_samples) - 1)]
    
    def stats(self) -> dict:
        """Get latency statistics."""
        return {
            "name": self.name,
            "count": len(self._samples),
            "avg_ms": round(self.avg, 2),
            "p50_ms": round(self.p50, 2),
            "p95_ms": round(self.p95, 2),
            "p99_ms": round(self.p99, 2),
        }


# Global latency trackers
_latency_trackers: dict[str, LatencyTracker] = {}


def get_latency_tracker(name: str) -> LatencyTracker:
    """Get or create a latency tracker."""
    if name not in _latency_trackers:
        _latency_trackers[name] = LatencyTracker(name=name)
    return _latency_trackers[name]


def timed(tracker_name: str):
    """Decorator to track function latency."""
    def decorator(func: Callable):
        tracker = get_latency_tracker(tracker_name)
        
        async def async_wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                return await func(*args, **kwargs)
            finally:
                duration_ms = (time.perf_counter() - start) * 1000
                tracker.record(duration_ms)
        
        def sync_wrapper(*args, **kwargs):
            start = time.perf_counter()
            try:
                return func(*args, **kwargs)
            finally:
                duration_ms = (time.perf_counter() - start) * 1000
                tracker.record(duration_ms)
        
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        return sync_wrapper
    
    return decorator


async def warm_up_connections(redis_factory: Callable) -> None:
    """
    Warm up connections on startup to minimize first-request latency.
    """
    _logger.info("warming_up_connections")
    
    try:
        # Warm up Redis connection
        timeout_s = float(os.getenv("REDIS_WARMUP_TIMEOUT_SECONDS", "3"))
        redis = await asyncio.wait_for(redis_factory(), timeout=timeout_s)
        await asyncio.wait_for(redis.ping(), timeout=timeout_s)
        _logger.info("redis_warmed_up")
    except Exception as e:
        _logger.warning("redis_warmup_failed", error=str(e))
    
    # Pre-load any cached configurations
    try:
        from app.auth.api_keys import _bootstrap_master_key
        _bootstrap_master_key()
        _logger.info("api_keys_warmed_up")
    except Exception as e:
        _logger.warning("api_keys_warmup_failed", error=str(e))
    
    _logger.info("warmup_complete")


def get_all_latency_stats() -> list[dict]:
    """Get statistics from all latency trackers."""
    return [tracker.stats() for tracker in _latency_trackers.values()]


# Optimized settings loader
@lru_cache(maxsize=1)
def get_optimized_settings() -> dict:
    """
    Load optimized settings for production.
    
    Cached to avoid repeated env lookups.
    """
    return {
        "worker_connections": int(os.getenv("WORKER_CONNECTIONS", "1000")),
        "keepalive_timeout": int(os.getenv("KEEPALIVE_TIMEOUT", "65")),
        "graceful_timeout": int(os.getenv("GRACEFUL_TIMEOUT", "30")),
        "max_requests": int(os.getenv("MAX_REQUESTS", "10000")),
        "max_requests_jitter": int(os.getenv("MAX_REQUESTS_JITTER", "1000")),
        "buffer_size": int(os.getenv("BUFFER_SIZE", "65536")),
        "backlog": int(os.getenv("BACKLOG", "2048")),
    }
