"""
Pydantic schemas for CAAE API request/response validation.
"""

from typing import Literal, Optional, Dict
from pydantic import BaseModel, Field


class QuantizationRequest(BaseModel):
    """Request payload for quantization analysis."""

    model_size_gb: float = Field(..., gt=0, description="Model size in GB")
    target_dtype: str = Field(
        ..., description="Target quantization dtype (fp16, bf16, int8, int4)"
    )
    memory_budget_gb: float = Field(
        ..., gt=0, description="Available memory budget in GB"
    )
    batch_size: int = Field(default=32, gt=0, description="Batch size")
    seq_len: int = Field(default=2048, gt=0, description="Sequence length")
    hidden_size: int = Field(default=4096, gt=0, description="Hidden dimension size")
    num_layers: int = Field(
        default=32, gt=0, description="Number of transformer layers"
    )


class QuantizationResponse(BaseModel):
    """Response payload for quantization analysis."""

    compression_ratio: float = Field(..., description="Compression ratio achieved")
    memory_savings_mb: float = Field(..., description="Memory savings in MB")
    latency_impact_ms: float = Field(..., description="Latency impact in milliseconds")
    accuracy_degradation: float = Field(
        ..., description="Expected accuracy degradation"
    )
    recommendation: str = Field(..., description="Recommended quantization strategy")
    kv_cache_compression: Dict = Field(..., description="KV cache compression analysis")


class EvictionRequest(BaseModel):
    """Request payload for eviction decision endpoint."""

    context_size_tokens: int = Field(
        ..., gt=0, description="Number of tokens in the context to potentially evict"
    )
    bandwidth_gbps: float = Field(
        default=16.0,
        gt=0,
        description="PCIe bandwidth in GB/s (Gen3=8, Gen4=16, Gen5=128)",
    )
    pcie_queue_depth_ms: float = Field(
        default=0.0,
        ge=0,
        description="Current PCIe queue depth/latency in milliseconds",
    )
    circuit_breaker_threshold_ms: float = Field(
        default=5.0, gt=0, description="Threshold for circuit breaker activation"
    )
    request_id: Optional[str] = Field(
        default=None, description="Optional request identifier for tracing"
    )
    tenant_id: Optional[str] = Field(
        default=None,
        description="Optional tenant identifier for multi-tenant scenarios",
    )
    layer_depth: Optional[int] = Field(
        default=None,
        ge=0,
        description="Layer depth of the block (0=early, higher=late)",
    )
    seq_index: Optional[int] = Field(
        default=None, ge=0, description="Sequence position index"
    )


class EvictionResponse(BaseModel):
    """Response payload from eviction decision endpoint."""

    action: Literal["swap", "recompute", "hold"] = Field(
        ..., description="Recommended eviction action"
    )
    predicted_swap_ms: float = Field(
        ..., description="Predicted cost of swap operation in milliseconds"
    )
    predicted_recompute_ms: float = Field(
        ..., description="Predicted cost of recomputation in milliseconds"
    )
    confidence: float = Field(
        ..., ge=0, le=1, description="Confidence score (0-1) for the decision"
    )
    circuit_breaker_active: bool = Field(
        ..., description="Whether circuit breaker is currently active"
    )
    request_id: Optional[str] = Field(
        default=None, description="Echo of request identifier for tracing"
    )
    reasoning: Optional[str] = Field(
        default=None, description="Human-readable explanation of decision"
    )


class HealthResponse(BaseModel):
    """Response payload for health check endpoint."""

    status: Literal["healthy", "degraded", "unhealthy"] = Field(
        ..., description="Service health status"
    )
    version: str = Field(..., description="Service version")
    uptime_seconds: float = Field(..., description="Seconds since service started")


class MetricsResponse(BaseModel):
    """Response payload for metrics endpoint."""

    total_decisions: int = Field(..., description="Total decisions made")
    swap_count: int = Field(..., description="Number of swap decisions")
    recompute_count: int = Field(..., description="Number of recompute decisions")
    hold_count: int = Field(..., description="Number of hold decisions")
    circuit_breaker_triggers: int = Field(
        ..., description="Circuit breaker activations"
    )
    avg_decision_latency_ms: float = Field(..., description="Average decision latency")
    p99_decision_latency_ms: float = Field(..., description="P99 decision latency")
