"""
Autotune Router Receipts — Tracks routing decisions and shadow lane outcomes.

This module provides receipt types for the autotune router:
- AutotuneRoutingReceipt: Classification output + provider selection
- ShadowLaneOutcomeReceipt: Per-lane execution result
- RoutingHistoryReceipt: Aggregated metrics for learning

These receipts enable:
- Auditable routing decisions
- Shadow lane comparison
- MLX classifier training data
- Cost/latency/accuracy tracking
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import Field

from receipts.base import BaseReceipt, Domain

logger = logging.getLogger("receipts.domains.autotune")


# ─────────────────────────────────────────────────────────────────────
# Enums
# ─────────────────────────────────────────────────────────────────────


class RoutingCapability(str, Enum):
    """Task capability for routing decisions."""
    REASONING = "reasoning"
    CODE = "code"
    VISION = "vision"
    FAST = "fast"
    QUALITY = "quality"
    BROKE = "broke"  # Free tier only


class RoutingLane(str, Enum):
    """Which lane the outcome came from."""
    CHAMPION = "champion"      # Current production
    SHADOW_A = "shadow_a"      # First shadow variant
    SHADOW_B = "shadow_b"      # Second shadow variant
    SHADOW_C = "shadow_c"      # Third shadow variant


class RoutingConfidence(str, Enum):
    """Confidence level of routing decision."""
    HIGH = "high"          # >0.8 confidence from classifier
    MEDIUM = "medium"      # 0.5-0.8 confidence
    LOW = "low"            # <0.5 confidence (fallback used)
    FALLBACK = "fallback"  # Classifier unavailable


# ─────────────────────────────────────────────────────────────────────
# Routing Features (input to classifier)
# ─────────────────────────────────────────────────────────────────────


class RoutingFeatures(BaseReceipt):
    """
    Features extracted from a request for routing decisions.

    These are the inputs to the MLX classifier.
    """
    receipt_type: str = Field(default="routing_features", description="Receipt type identifier")
    domain: str = Field(default="autotune", description="Domain namespace")

    # Request characteristics
    task_type: str = Field(description="Detected task type (e.g., 'code_review', 'reasoning', 'chat')")
    estimated_complexity: float = Field(ge=0.0, le=1.0, description="Estimated complexity 0-1")
    estimated_tokens: int = Field(ge=0, description="Estimated input tokens")

    # Context
    sensitivity: str = Field(default="internal", description="Data sensitivity level")
    urgency: float = Field(default=0.5, ge=0.0, le=1.0, description="Urgency 0-1")
    cost_budget: Optional[float] = Field(default=None, description="Max cost budget USD")
    latency_budget_ms: Optional[int] = Field(default=None, description="Max latency budget ms")

    # Historical context
    similar_task_count: int = Field(default=0, description="Count of similar past tasks")
    similar_task_success_rate: Optional[float] = Field(default=None, description="Success rate on similar tasks")

    # Feature vector (for ML)
    feature_vector: List[float] = Field(default_factory=list, description="Numeric feature vector for classifier")


# ─────────────────────────────────────────────────────────────────────
# Autotune Routing Receipt
# ─────────────────────────────────────────────────────────────────────


class AutotuneRoutingReceipt(BaseReceipt):
    """
    Records a routing decision made by the autotune router.

    Captures the classifier output, provider selection, and rationale.
    """
    receipt_type: str = Field(default="autotune_routing", description="Receipt type identifier")
    domain: str = Field(default="autotune", description="Domain namespace")

    # Request identification
    request_id: str = Field(description="ID of the request being routed")
    capability: RoutingCapability = Field(description="Requested capability")

    # Classifier output
    classifier_used: str = Field(default="mlx", description="Classifier type (mlx, heuristic, fallback)")
    classifier_latency_ms: float = Field(ge=0, description="Classifier latency in ms")
    classifier_confidence: float = Field(ge=0.0, le=1.0, description="Classification confidence")
    confidence_level: RoutingConfidence = Field(description="Binned confidence level")

    # Provider selection
    selected_provider: str = Field(description="Provider selected for this request")
    selected_model: str = Field(description="Model selected")
    fallback_chain: List[str] = Field(default_factory=list, description="Fallback providers if primary fails")

    # Alternative scores (for audit)
    provider_scores: Dict[str, float] = Field(default_factory=dict, description="Scores for each candidate provider")

    # Features used (reference)
    features_receipt_id: Optional[str] = Field(default=None, description="ID of RoutingFeatures receipt")

    # Cost/latency estimates
    estimated_cost_usd: Optional[float] = Field(default=None, description="Estimated cost for this provider")
    estimated_latency_ms: Optional[int] = Field(default=None, description="Estimated latency")

    # Shadow configuration
    shadow_lanes_enabled: bool = Field(default=False, description="Whether shadow lanes are active")
    shadow_providers: List[str] = Field(default_factory=list, description="Providers running in shadow")


# ─────────────────────────────────────────────────────────────────────
# Shadow Lane Outcome Receipt
# ─────────────────────────────────────────────────────────────────────


class ShadowLaneOutcomeReceipt(BaseReceipt):
    """
    Records the outcome from a shadow lane execution.

    One receipt per lane per request. Enables comparison across lanes.
    """
    receipt_type: str = Field(default="shadow_lane_outcome", description="Receipt type identifier")
    domain: str = Field(default="autotune", description="Domain namespace")

    # Linkage
    routing_receipt_id: str = Field(description="ID of the AutotuneRoutingReceipt")
    request_id: str = Field(description="ID of the original request")

    # Lane identification
    lane: RoutingLane = Field(description="Which lane this outcome is from")
    provider: str = Field(description="Provider used in this lane")
    model: str = Field(description="Model used")
    variant_id: str = Field(description="Variant identifier for shadow scoring")

    # Execution metrics
    latency_ms: float = Field(ge=0, description="Actual execution latency")
    cost_usd: float = Field(ge=0, default=0.0, description="Actual cost")
    tokens_input: int = Field(ge=0, default=0, description="Input tokens consumed")
    tokens_output: int = Field(ge=0, default=0, description="Output tokens generated")

    # Quality metrics
    success: bool = Field(description="Whether execution succeeded")
    error_code: Optional[str] = Field(default=None, description="Error code if failed")
    error_message: Optional[str] = Field(default=None, description="Error message if failed")

    # Prediction quality (if available)
    confidence: Optional[float] = Field(default=None, description="Model's confidence in response")
    quality_score: Optional[float] = Field(default=None, description="Quality score if evaluated")

    # Outcome linkage (for calibration)
    outcome_ref: Optional[str] = Field(default=None, description="Reference to outcome receipt when resolved")


# ─────────────────────────────────────────────────────────────────────
# Routing History Receipt (Aggregated)
# ─────────────────────────────────────────────────────────────────────


class RoutingHistoryReceipt(BaseReceipt):
    """
    Aggregated routing history for a provider+capability bucket.

    Used for:
    - Training the MLX classifier
    - Computing routing weights
    - Detecting drift
    """
    receipt_type: str = Field(default="routing_history", description="Receipt type identifier")
    domain: str = Field(default="autotune", description="Domain namespace")

    # Bucket identification
    provider: str = Field(description="Provider name")
    capability: RoutingCapability = Field(description="Capability")
    time_window: str = Field(description="Aggregation window (e.g., '7d', '30d')")

    # Volume metrics
    total_requests: int = Field(ge=0, description="Total requests in window")
    successful_requests: int = Field(ge=0, description="Successful requests")
    failed_requests: int = Field(ge=0, description="Failed requests")

    # Quality metrics
    avg_success_rate: float = Field(ge=0.0, le=1.0, description="Average success rate")
    avg_confidence: Optional[float] = Field(default=None, description="Average confidence")
    avg_quality_score: Optional[float] = Field(default=None, description="Average quality score")

    # Cost metrics
    total_cost_usd: float = Field(ge=0, default=0.0, description="Total cost in window")
    avg_cost_per_request: float = Field(ge=0, default=0.0, description="Average cost per request")

    # Latency metrics
    avg_latency_ms: float = Field(ge=0, description="Average latency")
    p50_latency_ms: float = Field(ge=0, description="P50 latency")
    p95_latency_ms: float = Field(ge=0, description="P95 latency")
    p99_latency_ms: float = Field(ge=0, description="P99 latency")

    # Computed weight (for routing)
    computed_weight: float = Field(ge=0.0, le=1.0, description="Computed routing weight")
    weight_factors: Dict[str, float] = Field(default_factory=dict, description="Factors contributing to weight")


# ─────────────────────────────────────────────────────────────────────
# Helper Functions
# ─────────────────────────────────────────────────────────────────────


def confidence_to_level(confidence: float) -> RoutingConfidence:
    """Convert numeric confidence to binned level."""
    if confidence >= 0.8:
        return RoutingConfidence.HIGH
    elif confidence >= 0.5:
        return RoutingConfidence.MEDIUM
    else:
        return RoutingConfidence.LOW


def create_routing_receipt(
    request_id: str,
    capability: RoutingCapability,
    selected_provider: str,
    selected_model: str,
    classifier_confidence: float,
    classifier_latency_ms: float,
    provider_scores: Dict[str, float],
    shadow_providers: List[str] = None,
    features_receipt_id: str = None,
    estimated_cost_usd: float = None,
    estimated_latency_ms: int = None,
) -> AutotuneRoutingReceipt:
    """
    Create an AutotuneRoutingReceipt with defaults.
    """
    return AutotuneRoutingReceipt(
        receipt_id=f"atr_{uuid.uuid4().hex[:12]}",
        request_id=request_id,
        capability=capability,
        classifier_used="mlx" if classifier_confidence > 0 else "fallback",
        classifier_latency_ms=classifier_latency_ms,
        classifier_confidence=classifier_confidence,
        confidence_level=confidence_to_level(classifier_confidence),
        selected_provider=selected_provider,
        selected_model=selected_model,
        fallback_chain=[],  # Populated from deployment profile
        provider_scores=provider_scores,
        features_receipt_id=features_receipt_id,
        estimated_cost_usd=estimated_cost_usd,
        estimated_latency_ms=estimated_latency_ms,
        shadow_lanes_enabled=bool(shadow_providers),
        shadow_providers=shadow_providers or [],
    )


def create_shadow_outcome(
    routing_receipt_id: str,
    request_id: str,
    lane: RoutingLane,
    provider: str,
    model: str,
    variant_id: str,
    latency_ms: float,
    cost_usd: float,
    success: bool,
    confidence: float = None,
    error_code: str = None,
    error_message: str = None,
) -> ShadowLaneOutcomeReceipt:
    """
    Create a ShadowLaneOutcomeReceipt with defaults.
    """
    return ShadowLaneOutcomeReceipt(
        receipt_id=f"slo_{uuid.uuid4().hex[:12]}",
        routing_receipt_id=routing_receipt_id,
        request_id=request_id,
        lane=lane,
        provider=provider,
        model=model,
        variant_id=variant_id,
        latency_ms=latency_ms,
        cost_usd=cost_usd,
        success=success,
        confidence=confidence,
        error_code=error_code,
        error_message=error_message,
    )


__all__ = [
    "RoutingCapability",
    "RoutingLane",
    "RoutingConfidence",
    "RoutingFeatures",
    "AutotuneRoutingReceipt",
    "ShadowLaneOutcomeReceipt",
    "RoutingHistoryReceipt",
    "confidence_to_level",
    "create_routing_receipt",
    "create_shadow_outcome",
]
