"""
vLLM Integration Layer for CAAE Advisor.

Wraps vLLM's BlockSpaceManager to transparently query CAAE for eviction decisions.
"""

import asyncio
import logging
import time
from collections import deque
from dataclasses import dataclass
from typing import Any, Optional, Deque

import httpx

try:
    import pynvml
    HAS_PYNVML = True
except ImportError:
    HAS_PYNVML = False
    logger.warning("pynvml not found. Real PCIe monitoring disabled.")

from .config import CAAEConfig, get_config

logger = logging.getLogger("caae-vllm")


@dataclass
class QueueDepthSample:
    """Single queue depth measurement."""

    timestamp: float
    queue_depth_ms: float
    context_tokens: int
    bandwidth_gbps: float


@dataclass
class QueueDepthStats:
    """Queue depth statistics for monitoring."""

    current_ms: float = 0.0
    min_ms: float = 0.0
    max_ms: float = 0.0
    avg_ms: float = 0.0
    p50_ms: float = 0.0
    p95_ms: float = 0.0
    p99_ms: float = 0.0
    total_measurements: int = 0
    threshold_exceeded_count: int = 0


class CAEBlockManagerWrapper:
    """
    Wraps vLLM's BlockSpaceManager to query CAAE before eviction.

    This wrapper intercepts eviction requests and queries the CAAE Advisor
    service to determine the optimal action. If CAAE is unavailable,
    it falls back to the original LRU policy gracefully.

    Includes queue depth measurement to track PCIe saturation and validate
    circuit breaker thresholds against real hardware behavior.

    Usage:
    ------
    from caae.vllm_integration import CAEBlockManagerWrapper

    # After creating the LLM engine
    engine = AsyncLLMEngine.from_engine_args(engine_args)

    # Wrap the block manager
    original_manager = engine.cache_engine.block_manager
    wrapped_manager = CAEBlockManagerWrapper(original_manager)
    engine.cache_engine.block_manager = wrapped_manager

    # Now all eviction decisions go through CAAE
    """

    # Queue depth measurement configuration
    QUEUE_DEPTH_SAMPLE_WINDOW = 300  # Keep 5 min of samples
    QUEUE_DEPTH_MEASUREMENT_INTERVAL_MS = 100  # Sample every 100ms

    def __init__(
        self,
        original_manager: Any,
        caae_advisor_url: str = "http://caae-advisor:8000/v1/decide",
        bandwidth_gbps: float = 16.0,
        timeout_ms: float = 5.0,
        config: CAAEConfig | None = None,
        enable_queue_depth_measurement: bool = True,
    ):
        """
        Initialize the CAAE-aware block manager wrapper.

        Args:
            original_manager: The original vLLM BlockSpaceManager to wrap
            caae_advisor_url: URL of the CAAE Advisor service
            bandwidth_gbps: PCIe bandwidth in GB/s (default: Gen 4)
            timeout_ms: HTTP timeout in milliseconds
            config: Optional CAAE configuration
            enable_queue_depth_measurement: Enable queue depth tracking
        """
        self.original = original_manager
        self.caae_url = caae_advisor_url
        self.bandwidth_gbps = bandwidth_gbps
        self.timeout_sec = timeout_ms / 1000.0
        self.config = config or get_config()
        self.enable_queue_depth_measurement = enable_queue_depth_measurement

        # Async HTTP client
        self._client: Optional[httpx.AsyncClient] = None

        # Queue depth tracking
        self._queue_depth_samples: Deque[QueueDepthSample] = deque(
            maxlen=self.QUEUE_DEPTH_SAMPLE_WINDOW
        )
        self._queue_depth_stats = QueueDepthStats()
        self._last_measurement_time = 0.0

        # Metrics
        self.decisions_made = 0
        self.swaps_executed = 0
        self.recomputes_triggered = 0
        self.holds_issued = 0
        self.fallbacks_triggered = 0
        self.caae_errors = 0

        # NVML Initialization for PCIe monitoring
        self._nvml_handle = None
        if HAS_PYNVML and self.enable_queue_depth_measurement:
            try:
                pynvml.nvmlInit()
                # Default to GPU 0, or find a way to pass this in
                self._nvml_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
                logger.info("NVML initialized for PCIe monitoring on GPU 0")
            except Exception as e:
                logger.error(f"Failed to initialize NVML: {e}")
                self.enable_queue_depth_measurement = False

        logger.info(
            f"CAEBlockManagerWrapper initialized: "
            f"url={caae_advisor_url}, bandwidth={bandwidth_gbps}GB/s, "
            f"queue_depth_measurement={enable_queue_depth_measurement}"
        )

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create HTTP client."""
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=self.timeout_sec)
        return self._client

    async def _query_caae(
        self,
        context_tokens: int,
        queue_depth_ms: float = 0.0,
        request_id: str | None = None,
    ) -> dict | None:
        """
        Query CAAE Advisor for eviction decision.

        Args:
            context_tokens: Number of tokens in context
            queue_depth_ms: Current PCIe queue depth
            request_id: Optional request ID for tracing

        Returns:
            Decision dict or None if query failed
        """
        try:
            client = await self._get_client()
            response = await client.post(
                self.caae_url,
                json={
                    "context_size_tokens": context_tokens,
                    "bandwidth_gbps": self.bandwidth_gbps,
                    "pcie_queue_depth_ms": queue_depth_ms,
                    "circuit_breaker_threshold_ms": self.config.circuit_breaker_threshold_ms,
                    "request_id": request_id,
                },
            )
            response.raise_for_status()
            return response.json()

        except httpx.TimeoutException:
            logger.warning(
                f"CAAE timeout after {self.timeout_sec*1000:.0f}ms, "
                f"falling back to baseline LRU"
            )
            self.fallbacks_triggered += 1
            return None

        except Exception as e:
            logger.error(f"CAAE query error: {e}")
            self.caae_errors += 1
            return None

    async def allocate_eviction_candidate(
        self,
        sequence: Any,
        num_tokens: int,
        reason: str = "memory_pressure",
    ) -> Optional[int]:
        """
        Query CAAE and decide whether to evict this block.

        This method intercepts eviction requests and queries CAAE
        to determine if eviction (swap) is optimal. Includes queue depth
        measurement for circuit breaker validation.

        Args:
            sequence: The vLLM Sequence object
            num_tokens: Number of tokens in the context
            reason: Reason for eviction

        Returns:
            Block ID to evict if action is "swap", None otherwise
        """
        # Measure queue depth
        if self.enable_queue_depth_measurement:
            queue_depth_ms = self._measure_queue_depth(num_tokens)
        else:
            queue_depth_ms = 0.0

        request_id = f"seq-{getattr(sequence, 'request_id', 'unknown')}"

        # Query CAAE
        decision = await self._query_caae(
            context_tokens=num_tokens,
            queue_depth_ms=queue_depth_ms,
            request_id=request_id,
        )

        self.decisions_made += 1

        if decision is None:
            # Fallback to original policy
            return self.original.allocate_eviction_candidate(
                sequence, num_tokens, reason
            )

        action = decision.get("action", "swap")

        if action == "swap":
            self.swaps_executed += 1
            logger.debug(
                f"CAAE: SWAP {num_tokens}t (qd={queue_depth_ms:.2f}ms) "
                f"(swap={decision['predicted_swap_ms']:.1f}ms < "
                f"recompute={decision['predicted_recompute_ms']:.1f}ms)"
            )
            # Proceed with original eviction
            return self.original.allocate_eviction_candidate(
                sequence, num_tokens, reason
            )

        elif action == "recompute":
            self.recomputes_triggered += 1
            logger.debug(
                f"CAAE: RECOMPUTE {num_tokens}t (qd={queue_depth_ms:.2f}ms) "
                f"(confidence={decision['confidence']:.1%})"
            )
            # Don't evict; signal to recompute instead
            return None

        else:  # "hold"
            self.holds_issued += 1
            logger.debug(f"CAAE: HOLD block (qd={queue_depth_ms:.2f}ms)")
            return None

    def _get_pcie_queue_depth(self) -> float:
        """
        Get current PCIe queue depth.

        Uses NVML to get PCIe throughput and estimates queueing delay
        using an M/M/1 queue model approximation.
        """
        if not self._nvml_handle or not HAS_PYNVML:
            return 0.0

        try:
            # Get TX and RX throughput in KB/s
            tx_kbps = pynvml.nvmlDeviceGetPcieThroughput(self._nvml_handle, pynvml.NVML_PCIE_UTIL_TX_BYTES)
            rx_kbps = pynvml.nvmlDeviceGetPcieThroughput(self._nvml_handle, pynvml.NVML_PCIE_UTIL_RX_BYTES)
            
            # Convert to GB/s
            total_gbps = (tx_kbps + rx_kbps) / 1024 / 1024
            
            # Calculate utilization (assuming bidirectional max theoretical)
            # Gen4 x16 is ~32GB/s per direction, so 64GB/s total theoretical max?
            # Actually usually bottlenecked by one direction. Let's take max of TX/RX relative to uni-directional limit.
            # 16 GB/s is the default passed in __init__, let's use that.
            
            max_gbps = self.bandwidth_gbps
            utilization = min(0.99, total_gbps / max_gbps) # Cap at 99% to avoid divide by zero
            
            # M/M/1 Queue Delay Approximation
            # Delay = Service_Time / (1 - Utilization)
            # Service_Time for a 4KB page at 16GB/s is ~0.25us. 
            # But overhead is higher. Let's assume base latency of 0.05ms (50us).
            base_latency_ms = 0.05
            estimated_delay_ms = base_latency_ms / (1.0 - utilization)
            
            # If utilization is very low, delay is just base latency
            # If utilization is high (e.g. 90%), delay is 0.05 / 0.1 = 0.5ms
            # If utilization is 99%, delay is 0.05 / 0.01 = 5ms (triggering threshold)
            
            return estimated_delay_ms

        except Exception as e:
            logger.warning(f"Failed to query PCIe stats: {e}")
            return 0.0

    def _measure_queue_depth(
        self,
        context_tokens: int,
    ) -> float:
        """
        Measure and record current queue depth.

        This method samples the PCIe queue depth at regular intervals
        and updates statistics for monitoring and circuit breaker validation.

        Args:
            context_tokens: Number of tokens in context (for correlation)

        Returns:
            Current queue depth in milliseconds
        """
        current_time = time.time()

        # Only sample at configured interval
        if (
            current_time - self._last_measurement_time
        ) * 1000 < self.QUEUE_DEPTH_MEASUREMENT_INTERVAL_MS:
            return self._queue_depth_stats.current_ms

        self._last_measurement_time = current_time
        queue_depth_ms = self._get_pcie_queue_depth()

        # Record sample
        sample = QueueDepthSample(
            timestamp=current_time,
            queue_depth_ms=queue_depth_ms,
            context_tokens=context_tokens,
            bandwidth_gbps=self.bandwidth_gbps,
        )
        self._queue_depth_samples.append(sample)

        # Update statistics
        self._update_queue_depth_stats(queue_depth_ms)

        return queue_depth_ms

    def _update_queue_depth_stats(self, queue_depth_ms: float) -> None:
        """
        Update queue depth statistics from samples.

        Args:
            queue_depth_ms: Latest measurement
        """
        self._queue_depth_stats.current_ms = queue_depth_ms
        self._queue_depth_stats.total_measurements += 1

        if queue_depth_ms > self.config.circuit_breaker_threshold_ms:
            self._queue_depth_stats.threshold_exceeded_count += 1

        if len(self._queue_depth_samples) > 0:
            depths = [s.queue_depth_ms for s in self._queue_depth_samples]

            self._queue_depth_stats.min_ms = min(depths)
            self._queue_depth_stats.max_ms = max(depths)
            self._queue_depth_stats.avg_ms = sum(depths) / len(depths)

            # Calculate percentiles
            sorted_depths = sorted(depths)
            n = len(sorted_depths)
            self._queue_depth_stats.p50_ms = sorted_depths[n // 2]
            self._queue_depth_stats.p95_ms = sorted_depths[int(n * 0.95)]
            self._queue_depth_stats.p99_ms = sorted_depths[int(n * 0.99)]

    def get_queue_depth_stats(self) -> QueueDepthStats:
        """
        Get current queue depth statistics for monitoring.

        Returns:
            QueueDepthStats with current measurements
        """
        return self._queue_depth_stats

    def get_queue_depth_samples(
        self, max_age_seconds: float = 60.0
    ) -> list[QueueDepthSample]:
        """
        Get recent queue depth samples for analysis.

        Args:
            max_age_seconds: Only return samples from last N seconds

        Returns:
            List of QueueDepthSample objects
        """
        cutoff_time = time.time() - max_age_seconds
        return [s for s in self._queue_depth_samples if s.timestamp >= cutoff_time]

    def get_metrics(self) -> dict:
        """Get integration metrics for monitoring."""
        total = max(1, self.decisions_made)
        queue_stats = self.get_queue_depth_stats()

        return {
            "decisions_made": self.decisions_made,
            "swaps_executed": self.swaps_executed,
            "recomputes_triggered": self.recomputes_triggered,
            "holds_issued": self.holds_issued,
            "fallbacks_triggered": self.fallbacks_triggered,
            "caae_errors": self.caae_errors,
            "swap_ratio": round(self.swaps_executed / total, 3),
            "fallback_ratio": round(self.fallbacks_triggered / total, 3),
            # Queue depth metrics
            "queue_depth_current_ms": round(queue_stats.current_ms, 3),
            "queue_depth_min_ms": round(queue_stats.min_ms, 3),
            "queue_depth_max_ms": round(queue_stats.max_ms, 3),
            "queue_depth_avg_ms": round(queue_stats.avg_ms, 3),
            "queue_depth_p50_ms": round(queue_stats.p50_ms, 3),
            "queue_depth_p95_ms": round(queue_stats.p95_ms, 3),
            "queue_depth_p99_ms": round(queue_stats.p99_ms, 3),
            "queue_depth_threshold_exceeded": queue_stats.threshold_exceeded_count,
        }

    async def close(self) -> None:
        """Clean up HTTP client and NVML."""
        if self._client:
            await self._client.aclose()
            self._client = None
            
        if self._nvml_handle and HAS_PYNVML:
            try:
                pynvml.nvmlShutdown()
            except:
                pass
            self._nvml_handle = None

    # Proxy all other methods to original manager
    def __getattr__(self, name: str) -> Any:
        """Proxy all other attribute access to original manager."""
        return getattr(self.original, name)


class SyncCAEBlockManagerWrapper:
    """
    Synchronous wrapper for non-async vLLM environments.

    Uses asyncio.run() to execute async CAAE queries.
    """

    def __init__(self, *args, **kwargs):
        self._async_wrapper = CAEBlockManagerWrapper(*args, **kwargs)

    def allocate_eviction_candidate(self, *args, **kwargs) -> Optional[int]:
        """Synchronous wrapper for eviction candidate allocation."""
        return asyncio.run(
            self._async_wrapper.allocate_eviction_candidate(*args, **kwargs)
        )

    def get_metrics(self) -> dict:
        return self._async_wrapper.get_metrics()

    def __getattr__(self, name: str) -> Any:
        return getattr(self._async_wrapper.original, name)
