#!/usr/bin/env python3
"""
CAAE vLLM Plugin - Drop-in Extension for LLM Inference Optimization

Integrates CAAE optimizations (Exp 8, 9, 10, 13) directly into vLLM's
request processing pipeline with minimal code changes.

Usage:
    # In your vLLM initialization:
    from caae.vllm_plugin import CAAAEPlugin

    plugin = CAAAEPlugin(config_path="caae_config.yaml")
    llm = LLM(model="meta-llama/Llama-2-70b-hf", plugins=[plugin])

Author: CAAE Research
Date: January 23, 2026
"""

import time
import logging
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from datetime import datetime
import threading
import json
from pathlib import Path

from .hybrid_optimizer import UserTier, RequestContext, get_optimizer
from .config import detect_gpu_memory_gb, get_pcie_bandwidth_with_fallback

logger = logging.getLogger(__name__)


# ============================================================================
# CONFIGURATION
# ============================================================================


@dataclass
class CAAAEConfig:
    """CAAE plugin configuration with auto-detection support."""

    # Optimization toggles
    enable_speculative_decoding_opt: bool = False  # Exp 8 - Conservative default
    enable_multi_gpu_coordination: bool = False  # Exp 9 - Requires 2+ GPUs
    enable_shared_kv_pooling: bool = False  # Exp 7 - Conservative default
    enable_adaptive_sla: bool = True  # Exp 10 - Safe to enable

    # SLA targets
    target_p95_ms: float = 300.0
    target_p99_ms: float = 500.0
    sla_strict_mode: bool = False

    # Cost model (with auto-detection)
    gpu_cost_per_hour: float = 1.5
    gpu_memory_gb: Optional[float] = None  # Auto-detected if None
    pcie_bandwidth_gbps: Optional[float] = None  # Auto-detected if None

    # Performance tuning
    prediction_accuracy_target: float = 0.71  # 71% for Exp 8
    circuit_breaker_timeout_ms: float = 5.0

    # Monitoring
    enable_telemetry: bool = True
    telemetry_sample_rate: float = 0.1  # 10% of requests
    metrics_export_interval_sec: int = 60

    # Fallback behavior
    fallback_on_error: bool = True  # Disable optimization if error
    strict_mode: bool = False  # Require success or fail request

    # Advisor configuration
    advisor_url: Optional[str] = None
    advisor_timeout_ms: float = 5.0
    advisor_check_on_startup: bool = True

    @classmethod
    def find_config_file(cls) -> Optional[Path]:
        """
        Search for config file in standard locations.
        
        Search order:
        1. ./caae_config.yaml (current directory)
        2. ~/.kvict/caae_config.yaml (user config)
        3. <package>/default_config.yaml (bundled defaults)
        
        Returns:
            Path to config file, or None if using bundled defaults.
        """
        # Check current directory
        cwd_config = Path("caae_config.yaml")
        if cwd_config.exists():
            logger.info(f"Using config from current directory: {cwd_config}")
            return cwd_config
        
        # Check user config directory
        user_config = Path.home() / ".kvict" / "caae_config.yaml"
        if user_config.exists():
            logger.info(f"Using user config: {user_config}")
            return user_config
        
        # Use bundled default
        default_config = Path(__file__).parent / "default_config.yaml"
        if default_config.exists():
            logger.info(f"Using bundled default config: {default_config}")
            return default_config
        
        logger.warning("No config file found, using hardcoded defaults")
        return None

    @classmethod
    def from_file(cls, path: Optional[str] = None) -> "CAAAEConfig":
        """
        Load config from YAML file with auto-detection.
        
        Args:
            path: Explicit config path, or None to search standard locations.
        
        Returns:
            CAAAEConfig instance with auto-detected hardware values.
        """
        import yaml
        
        # Find config file
        if path:
            config_path = Path(path)
            if not config_path.exists():
                logger.error(f"Config file not found: {path}")
                raise FileNotFoundError(f"Config file not found: {path}")
        else:
            config_path = cls.find_config_file()
        
        # Load YAML
        config_dict = {}
        if config_path:
            try:
                with open(config_path) as f:
                    raw_config = yaml.safe_load(f) or {}
                
                # Support both nested and flat formats
                config_dict = cls._normalize_config(raw_config)
                logger.debug(f"Loaded config from {config_path}")
            except Exception as e:
                logger.warning(f"Failed to load config from {config_path}: {e}")
        
        # Create instance with loaded values
        instance = cls(**{k: v for k, v in config_dict.items() if hasattr(cls, k)})
        
        # Apply auto-detection for missing values
        instance._auto_detect_hardware()
        
        return instance

    @staticmethod
    def _normalize_config(raw_config: Dict[str, Any]) -> Dict[str, Any]:
        """
        Normalize config from nested format to flat format.
        
        Supports both:
        - Nested: optimizations.adaptive_sla → enable_adaptive_sla
        - Flat: enable_adaptive_sla → enable_adaptive_sla
        """
        normalized = {}
        
        # Extract nested optimizations
        if "optimizations" in raw_config:
            opts = raw_config["optimizations"]
            normalized["enable_speculative_decoding_opt"] = opts.get("speculative_decoding", False)
            normalized["enable_multi_gpu_coordination"] = opts.get("multi_gpu_coordination", False)
            normalized["enable_shared_kv_pooling"] = opts.get("shared_kv_pooling", False)
            normalized["enable_adaptive_sla"] = opts.get("adaptive_sla", True)
        
        # Extract SLA config
        if "sla" in raw_config:
            sla = raw_config["sla"]
            normalized["target_p95_ms"] = sla.get("target_p95_ms", 300.0)
            normalized["target_p99_ms"] = sla.get("target_p99_ms", 500.0)
            normalized["sla_strict_mode"] = sla.get("strict_mode", False)
        
        # Extract cost model
        if "cost_model" in raw_config:
            cost = raw_config["cost_model"]
            normalized["gpu_cost_per_hour"] = cost.get("gpu_cost_per_hour", 1.5)
            
            # Handle "auto" string for auto-detection
            gpu_mem = cost.get("gpu_memory_gb")
            if gpu_mem != "auto" and gpu_mem is not None:
                normalized["gpu_memory_gb"] = float(gpu_mem)
            
            pcie_bw = cost.get("pcie_bandwidth_gbps")
            if pcie_bw != "auto" and pcie_bw is not None:
                normalized["pcie_bandwidth_gbps"] = float(pcie_bw)
        
        # Extract performance config
        if "performance" in raw_config:
            perf = raw_config["performance"]
            normalized["prediction_accuracy_target"] = perf.get("prediction_accuracy_target", 0.71)
            normalized["circuit_breaker_timeout_ms"] = perf.get("circuit_breaker_timeout_ms", 5.0)
        
        # Extract telemetry config
        if "telemetry" in raw_config:
            telem = raw_config["telemetry"]
            normalized["enable_telemetry"] = telem.get("enable", True)
            normalized["telemetry_sample_rate"] = telem.get("sample_rate", 0.1)
            normalized["metrics_export_interval_sec"] = telem.get("export_interval_sec", 60)
        
        # Extract fallback config
        if "fallback" in raw_config:
            fallback = raw_config["fallback"]
            normalized["fallback_on_error"] = fallback.get("on_error", True)
            normalized["strict_mode"] = fallback.get("strict_mode", False)
        
        # Extract advisor config
        if "advisor" in raw_config:
            advisor = raw_config["advisor"]
            normalized["advisor_url"] = advisor.get("url")
            normalized["advisor_timeout_ms"] = advisor.get("timeout_ms", 5.0)
            normalized["advisor_check_on_startup"] = advisor.get("check_on_startup", True)
        
        # Also merge flat keys (backward compatibility)
        flat_keys = [
            "enable_speculative_decoding_opt", "enable_multi_gpu_coordination", 
            "enable_shared_kv_pooling", "enable_adaptive_sla",
            "prediction_accuracy_target", "circuit_breaker_timeout_ms", 
            "sla_strict_mode", "enable_telemetry", "telemetry_sample_rate", 
            "metrics_export_interval_sec", "fallback_on_error", "strict_mode"
        ]
        for key in flat_keys:
            if key in raw_config:
                normalized[key] = raw_config[key]
        
        return normalized

    def _auto_detect_hardware(self) -> None:
        """Auto-detect hardware specifications if not configured."""
        # Auto-detect GPU memory
        if self.gpu_memory_gb is None:
            detected_mem = detect_gpu_memory_gb()
            if detected_mem:
                self.gpu_memory_gb = detected_mem
            else:
                self.gpu_memory_gb = 80.0  # Conservative default (A100)
                logger.warning(
                    f"Could not detect GPU memory, using default: {self.gpu_memory_gb} GB. "
                    "Install torch with CUDA support for accurate detection."
                )
        
        # Auto-detect PCIe bandwidth
        if self.pcie_bandwidth_gbps is None:
            self.pcie_bandwidth_gbps = get_pcie_bandwidth_with_fallback()


# ============================================================================
# METRICS & TELEMETRY
# ============================================================================


@dataclass
class RequestMetrics:
    """Per-request metrics."""

    request_id: str
    timestamp: float = field(default_factory=time.time)
    model: str = ""

    # Baseline metrics
    baseline_latency_ms: Optional[float] = None
    baseline_cache_hit_rate: float = 0.0

    # CAAE metrics
    caae_latency_ms: Optional[float] = None
    caae_cache_hit_rate: float = 0.0
    improvement_percent: float = 0.0

    # Experiment metrics
    exp8_prediction_accuracy: Optional[float] = None
    exp9_nvlink_util: Optional[float] = None
    exp10_sla_met: bool = False

    # Cost
    gpu_cost_usd: float = 0.0
    estimated_saving_usd: float = 0.0

    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            "request_id": self.request_id,
            "timestamp": datetime.fromtimestamp(self.timestamp).isoformat(),
            "model": self.model,
            "baseline_latency_ms": self.baseline_latency_ms,
            "caae_latency_ms": self.caae_latency_ms,
            "improvement_percent": self.improvement_percent,
            "exp8_prediction_accuracy": self.exp8_prediction_accuracy,
            "exp9_nvlink_util": self.exp9_nvlink_util,
            "exp10_sla_met": self.exp10_sla_met,
            "estimated_saving_usd": self.estimated_saving_usd,
        }


@dataclass
class AggregatedMetrics:
    """Aggregated metrics for dashboard."""

    window_start: float = field(default_factory=time.time)
    window_end: float = field(default_factory=time.time)

    # Throughput
    total_requests: int = 0
    total_tokens: int = 0
    qps: float = 0.0  # Requests per second
    tokens_per_second: float = 0.0

    # Latency
    p50_latency_ms: float = 0.0
    p99_latency_ms: float = 0.0
    p999_latency_ms: float = 0.0
    mean_latency_ms: float = 0.0

    # Improvement
    avg_improvement_percent: float = 0.0
    total_cost_saved_usd: float = 0.0

    # SLA compliance
    sla_compliance_rate: float = 0.0

    # GPU utilization
    avg_gpu_util_percent: float = 0.0
    avg_memory_util_percent: float = 0.0

    def to_dict(self) -> Dict:
        """Convert to dictionary."""
        return {
            "window_start": datetime.fromtimestamp(self.window_start).isoformat(),
            "window_end": datetime.fromtimestamp(self.window_end).isoformat(),
            "total_requests": self.total_requests,
            "total_tokens": self.total_tokens,
            "qps": round(self.qps, 2),
            "tokens_per_second": round(self.tokens_per_second, 0),
            "p50_latency_ms": round(self.p50_latency_ms, 2),
            "p99_latency_ms": round(self.p99_latency_ms, 2),
            "p999_latency_ms": round(self.p999_latency_ms, 2),
            "mean_latency_ms": round(self.mean_latency_ms, 2),
            "avg_improvement_percent": round(self.avg_improvement_percent, 1),
            "total_cost_saved_usd": round(self.total_cost_saved_usd, 2),
            "sla_compliance_rate": round(self.sla_compliance_rate * 100, 1),
            "avg_gpu_util_percent": round(self.avg_gpu_util_percent, 1),
            "avg_memory_util_percent": round(self.avg_memory_util_percent, 1),
        }


# ============================================================================
# PLUGIN IMPLEMENTATION
# ============================================================================


class CAAAEPlugin:
    """
    vLLM Plugin for CAAE Optimizations (Zero-Config).

    Automatically detects GPU specs and applies validated optimizations
    from Experiments 7-13.

    **Quick Start** (Zero Configuration):
        >>> from caae.vllm_plugin import CAAAEPlugin
        >>> from vllm import LLM
        >>> 
        >>> plugin = CAAAEPlugin()  # Auto-detects GPU, uses sensible defaults
        >>> llm = LLM(model="meta-llama/Llama-2-70b-hf", plugins=[plugin])

    **Custom Configuration**:
        >>> plugin = CAAAEPlugin(config_path="my_config.yaml")
        >>> llm = LLM(model="...", plugins=[plugin])

    **One-Liner** (recommended):
        >>> from caae.vllm_plugin import create_optimized_llm
        >>> llm = create_optimized_llm("meta-llama/Llama-2-70b-hf")
    """

    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize CAAE plugin with zero-config support.
        
        Args:
            config_path: Path to config file, or None for auto-detection.
                        When None, searches: ./caae_config.yaml → ~/.kvict/caae_config.yaml → bundled defaults
        """
        # Load config with auto-detection
        self.config = CAAAEConfig.from_file(config_path)
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

        # Metrics storage
        self.request_metrics: Dict[str, RequestMetrics] = {}
        self.aggregated_metrics = AggregatedMetrics()
        self.metrics_lock = threading.Lock()

        # Request tracking
        self.request_count = 0
        self.baseline_latencies: List[float] = []
        self.caae_latencies: List[float] = []

        # Hybrid optimizer
        optimizer_config = {
            "enable_caching": True,
            "enable_compression": True,
            "enable_adaptive_prompting": True,
            "enable_priority_allocation": True,
            "common_query_threshold": 3,
        }
        self.hybrid_optimizer = get_optimizer(optimizer_config)

        self.logger.info("CAAE Plugin initialized with Hybrid Optimizer")
        self.logger.info(f"Config: {self.config}")

    # ========================================================================
    # vLLM INTEGRATION HOOKS
    # ========================================================================

    def on_request_start(self, request: Dict[str, Any]) -> Dict[str, Any]:
        """
        Hook: Called when request arrives (before processing).

        Experiment 10: Adaptive SLA prioritization
        - Extract SLA deadline from request
        - Compute adaptive eviction threshold
        - Extract tenant_id for multi-tenant scenarios

        Hybrid Optimizer: optimization
        - Apply hybrid strategy for contradiction reduction
        - Optimize token usage based on user tier
        - Enable predictive caching
        """
        # Ensure a stable request_id for downstream hooks/metrics
        if request.get("request_id"):
            request_id = str(request["request_id"])
        else:
            with self.metrics_lock:
                request_id = f"req_{self.request_count}"
                self.request_count += 1
            request["request_id"] = request_id

        # Extract tenant_id from request (multi-tenant support)
        tenant_id = self._extract_tenant_id(request)
        if tenant_id:
            request["caae_tenant_id"] = tenant_id
            self.logger.debug(f"Req {request_id}: tenant_id={tenant_id}")

        # Extract user tier for optimization
        user_tier_str = request.get("user_tier") or request.get("tier") or "free"
        try:
            user_tier = UserTier(user_tier_str.lower())
        except ValueError:
            user_tier = UserTier.FREE  # Default to free tier

        request["caae_user_tier"] = user_tier.value

        try:
            if self.config.enable_adaptive_sla:
                # Exp 10: Adaptive SLA
                sla_deadline_ms = request.get("sla_deadline_ms", 3000)
                priority = self._infer_priority(sla_deadline_ms)

                # Compute adaptive threshold
                adaptive_threshold = self._compute_adaptive_threshold(sla_deadline_ms)

                request["caae_priority"] = priority  # Now returns float
                request["caae_adaptive_threshold"] = adaptive_threshold
                request["caae_sla_deadline"] = sla_deadline_ms
                request["caae_start_time"] = time.time()

                self.logger.debug(
                    f"Req {request_id}: adaptive SLA threshold {adaptive_threshold}ms "
                    f"for deadline {sla_deadline_ms}ms"
                )

            # Hybrid Optimizer: Create request context
            prompt = request.get("prompt", "") or request.get("messages", [{}])[-1].get(
                "content", ""
            )
            request_context = RequestContext(
                request_id=request_id,
                prompt=prompt,
                user_tier=user_tier,
                priority=request.get("caae_priority", 1.0),
                value_score=request.get("value_score", 1.0),
                difficulty=request.get("difficulty", "medium"),
                task_type=request.get("task_type", "general"),
                remaining_budget=request.get("remaining_budget", 10000),
                budget_percent_remaining=request.get("budget_percent_remaining", 1.0),
            )

            # Apply hybrid optimization
            opt_result = self.hybrid_optimizer.optimize_request(request_context)
            request["caae_optimization"] = {
                "served": opt_result.served,
                "tokens_used": opt_result.tokens_used,
                "quality_delivered": opt_result.quality_delivered,
                "cache_hit": opt_result.cache_hit,
                "strategy_used": opt_result.strategy_used,
                "contradiction_avoided": opt_result.contradiction_avoided,
            }

            if opt_result.cache_hit:
                self.logger.debug(
                    f"Req {request_id}: Cache hit! Tokens saved: {opt_result.tokens_used}"
                )

        except Exception as e:
            self.logger.error(f"Error in on_request_start: {e}")
            if self.config.fallback_on_error:
                return request
            raise

        return request

    def _extract_tenant_id(self, request: Dict[str, Any]) -> Optional[str]:
        """
        Extract tenant_id from a vLLM request.

        Looks for tenant_id in common locations:
        - request["tenant_id"]
        - request["tenant"]
        - request["user_id"]
        - request["metadata"]["tenant_id"]
        - request["params"]["tenant_id"]

        Args:
            request: vLLM request dictionary

        Returns:
            Tenant ID if found, None otherwise
        """
        # Check direct fields
        tenant_id = (
            request.get("tenant_id") or request.get("tenant") or request.get("user_id")
        )
        if tenant_id:
            return str(tenant_id)

        # Check in metadata
        if "metadata" in request:
            metadata = request["metadata"]
            if isinstance(metadata, dict):
                tenant_id = metadata.get("tenant_id") or metadata.get("tenant")
                if tenant_id:
                    return str(tenant_id)

        # Check in params
        if "params" in request:
            params = request["params"]
            if isinstance(params, dict):
                tenant_id = params.get("tenant_id") or params.get("tenant")
                if tenant_id:
                    return str(tenant_id)

        return None

    def on_cache_allocation(
        self, request: Dict[str, Any], cache_blocks: List[Dict]
    ) -> List[Dict]:
        """
        Hook: Called when allocating KV cache (before allocation).

        Experiment 8: Speculative decoding optimization
        - Skip caching low-confidence draft tokens

        Experiment 7: Shared KV pooling
        - Detect shared content and pool blocks
        """
        # Use the same request_id assigned at start; if missing, generate a unique one
        if request.get("request_id"):
            request_id = str(request["request_id"])
        else:
            with self.metrics_lock:
                request_id = f"req_{self.request_count}"
                self.request_count += 1
            request["request_id"] = request_id

        try:
            if self.config.enable_shared_kv_pooling:
                # Exp 7: Shared pooling
                cache_blocks = self._apply_shared_pooling(request_id, cache_blocks)

            if self.config.enable_speculative_decoding_opt:
                # Exp 8: Selective caching
                cache_blocks = self._apply_selective_caching(request_id, cache_blocks)

            return cache_blocks

        except Exception as e:
            self.logger.error(f"Error in on_cache_allocation: {e}")
            if self.config.fallback_on_error:
                return cache_blocks
            raise

    def on_request_end(self, request: Dict[str, Any], response: Dict[str, Any]) -> None:
        """
        Hook: Called when request completes (after processing).

        Collect metrics for dashboard and billing.
        """
        request_id = request.get("request_id", "unknown")

        try:
            start_time = request.get("caae_start_time", time.time())
            latency_ms = (time.time() - start_time) * 1000

            metrics = RequestMetrics(
                request_id=request_id,
                model=request.get("model", "unknown"),
                caae_latency_ms=latency_ms,
                exp10_sla_met=self._check_sla_compliance(request, latency_ms),
            )

            # Store metrics
            with self.metrics_lock:
                self.request_metrics[request_id] = metrics
                self.caae_latencies.append(latency_ms)

                # Limit memory: keep last 10k requests
                if len(self.request_metrics) > 10000:
                    oldest_key = min(self.request_metrics.keys())
                    del self.request_metrics[oldest_key]

        except Exception as e:
            self.logger.error(f"Error in on_request_end: {e}")

    # ========================================================================
    # EXPERIMENT IMPLEMENTATIONS
    # ========================================================================

    def _apply_selective_caching(
        self, request_id: str, cache_blocks: List[Dict]
    ) -> List[Dict]:
        """
        Experiment 8: Selective caching for speculative decoding.

        Only cache draft tokens predicted to be accepted (71.4% accuracy).
        Skip low-confidence tokens to save cache space.
        """
        if not self.config.enable_speculative_decoding_opt:
            return cache_blocks

        filtered_blocks = []

        for block in cache_blocks:
            # Check if this is a draft token
            if block.get("is_draft_token", False):
                # Prediction confidence from draft model
                confidence = block.get("prediction_confidence", 0.5)

                # Only cache if confidence > threshold (71.4% model)
                # Use conservative threshold to match 71.4% accuracy target
                if confidence > 0.65:
                    filtered_blocks.append(block)
                else:
                    # Skip caching this draft token
                    self.logger.debug(
                        f"Exp8: Skipped draft token (conf={confidence:.2f})"
                    )
            else:
                # Always cache non-draft tokens
                filtered_blocks.append(block)

        # Metrics
        cache_reduction_percent = (
            1 - len(filtered_blocks) / max(len(cache_blocks), 1)
        ) * 100
        if cache_reduction_percent > 0:
            self.logger.debug(
                f"Exp8: {cache_reduction_percent:.1f}% cache reduction via selective caching"
            )

        return filtered_blocks

    def _apply_shared_pooling(
        self, request_id: str, cache_blocks: List[Dict]
    ) -> List[Dict]:
        """
        Experiment 7: Shared KV cache pooling for RAG workloads.

        Detect shared content (system prompts, documents) and reuse blocks
        across concurrent requests instead of duplicating.
        """
        if not self.config.enable_shared_kv_pooling:
            return cache_blocks

        # In real implementation, would:
        # 1. Hash block content
        # 2. Check if hash exists in global pool
        # 3. Reference existing block instead of allocating new one
        # 4. Increment reference count

        # For MVP: just mark eligible blocks for pooling
        for block in cache_blocks:
            # Blocks with high content_hash matching potential
            if block.get("is_context_block", False):
                block["caae_poolable"] = True

        return cache_blocks

    def _compute_adaptive_threshold(self, sla_deadline_ms: float) -> float:
        """
        Experiment 10: Compute adaptive eviction threshold based on SLA.

        Algorithm:
        - If deadline < 500ms: aggressive (5ms threshold)
        - If deadline 500-1000ms: balanced (10ms threshold)
        - If deadline > 1000ms: conservative (15ms threshold)
        """
        if sla_deadline_ms < 500:
            return 5.0  # Aggressive
        elif sla_deadline_ms < 1000:
            return 10.0  # Balanced
        else:
            return 15.0  # Conservative

    def _infer_priority(self, sla_deadline_ms: float) -> float:
        """Infer request priority from SLA deadline (0.0 to 1.0)."""
        if sla_deadline_ms < 500:
            return 1.0  # Critical priority
        elif sla_deadline_ms < 1000:
            return 0.7  # Normal priority
        else:
            return 0.3  # Batch priority

    def _check_sla_compliance(self, request: Dict[str, Any], latency_ms: float) -> bool:
        """Check if request met its SLA deadline."""
        sla_deadline = request.get("caae_sla_deadline", 3000)
        return latency_ms <= sla_deadline

    # ========================================================================
    # METRICS & MONITORING
    # ========================================================================

    def get_metrics(self) -> AggregatedMetrics:
        """Get aggregated metrics for dashboard."""
        with self.metrics_lock:
            if not self.caae_latencies:
                return self.aggregated_metrics

            # Calculate percentiles
            sorted_latencies = sorted(self.caae_latencies)
            n = len(sorted_latencies)

            metrics = AggregatedMetrics(
                total_requests=len(self.request_metrics),
                qps=len(self.caae_latencies) / 60.0 if self.caae_latencies else 0,
                p50_latency_ms=sorted_latencies[int(n * 0.50)] if n > 0 else 0,
                p99_latency_ms=sorted_latencies[int(n * 0.99)] if n > 0 else 0,
                p999_latency_ms=sorted_latencies[int(n * 0.999)] if n > 0 else 0,
                mean_latency_ms=(
                    sum(self.caae_latencies) / len(self.caae_latencies)
                    if self.caae_latencies
                    else 0
                ),
                sla_compliance_rate=(
                    sum(1 for m in self.request_metrics.values() if m.exp10_sla_met)
                    / len(self.request_metrics)
                    if self.request_metrics
                    else 0
                ),
            )

            return metrics

    def get_hybrid_optimizer_metrics(self) -> Dict[str, Any]:
        """Get hybrid optimizer metrics for dashboard."""
        return self.hybrid_optimizer.get_metrics()

    def export_metrics_json(self, output_file: str) -> None:
        """Export metrics to JSON file for external monitoring."""
        with self.metrics_lock:
            data = {
                "timestamp": datetime.utcnow().isoformat() + "Z",
                "config": {
                    "enable_exp8": self.config.enable_speculative_decoding_opt,
                    "enable_exp9": self.config.enable_multi_gpu_coordination,
                    "enable_exp10": self.config.enable_adaptive_sla,
                    "enable_exp7": self.config.enable_shared_kv_pooling,
                },
                "aggregated": self.get_metrics().to_dict(),
                "hybrid_optimizer": self.get_hybrid_optimizer_metrics(),
                "recent_requests": [
                    m.to_dict() for m in list(self.request_metrics.values())[-100:]
                ],
            }

            Path(output_file).parent.mkdir(parents=True, exist_ok=True)
            with open(output_file, "w") as f:
                json.dump(data, f, indent=2)

            self.logger.info(f"Metrics exported to {output_file}")


# ============================================================================
# CONVENIENCE FUNCTIONS
# ============================================================================


def create_caae_plugin(config_path: Optional[str] = None) -> CAAAEPlugin:
    """
    Factory function to create CAAE plugin with defaults.

    Usage:
        from caae.vllm_plugin import create_caae_plugin

        plugin = create_caae_plugin()
        llm = LLM(model="...", plugins=[plugin])
    """
    return CAAAEPlugin(config_path=config_path)


def create_optimized_llm(model: str, config_path: Optional[str] = None, **vllm_kwargs):
    """
    One-line vLLM + CAAE setup with auto-detection (recommended).
    
    This is the easiest way to use CAAE with vLLM. It handles plugin
    creation, hardware auto-detection, and sensible defaults automatically.
    
    Args:
        model: HuggingFace model ID (e.g., "meta-llama/Llama-2-70b-hf")
        config_path: Optional path to CAAE config file (auto-detects if None)
        **vllm_kwargs: Additional arguments passed to vLLM LLM constructor
    
    Returns:
        vLLM LLM instance with CAAE optimizations enabled
    
    Example:
        >>> from caae.vllm_plugin import create_optimized_llm
        >>> 
        >>> # Zero-config quick start
        >>> llm = create_optimized_llm("meta-llama/Llama-2-70b-hf")
        >>> 
        >>> # With custom config
        >>> llm = create_optimized_llm(
        ...     "meta-llama/Llama-2-70b-hf",
        ...     config_path="production_config.yaml",
        ...     tensor_parallel_size=2
        ... )
        >>> 
        >>> # Generate
        >>> outputs = llm.generate(["Hello, my name is"], max_tokens=50)
    """
    try:
        from vllm import LLM
    except ImportError:
        raise ImportError(
            "vLLM is not installed. Install with: pip install vllm>=0.6.0"
        )
    
    plugin = CAAAEPlugin(config_path=config_path)
    logger.info(f"Creating optimized LLM for model: {model}")
    logger.info(f"Auto-detected: {plugin.config.gpu_memory_gb:.1f} GB GPU, "
                f"{plugin.config.pcie_bandwidth_gbps:.1f} GB/s PCIe")
    
    return LLM(model=model, plugins=[plugin], **vllm_kwargs)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # Test plugin
    plugin = CAAAEPlugin()

    # Simulate a request
    request = {
        "request_id": "test_001",
        "model": "meta-llama/Llama-2-70b",
        "sla_deadline_ms": 1000,
    }

    print("Testing CAAE vLLM Plugin...")
    print(f"Request: {request}")

    # Process request
    request = plugin.on_request_start(request)
    print(f"After on_request_start: {request}")

    # Simulate completion
    time.sleep(0.1)
    plugin.on_request_end(request, {})

    # Get metrics
    metrics = plugin.get_metrics()
    print(f"Metrics: {metrics.to_dict()}")
