#!/usr/bin/env python3
"""
Pydantic schemas for FlowTorch production API.
Typed contracts with validation for multi-state basin control.
"""

from typing import Optional, Literal, List, Dict, Any
from pydantic import BaseModel, Field, validator
from datetime import datetime, timezone


class QuantOverlay(BaseModel):
    """Quantum/phase modulation parameters for basin control."""

    # Phase modulation strength (0=off, 1=max)
    overlay_alpha: float = Field(0.5, ge=0.0, le=1.0, description="Phase modulation strength")

    # Frequency control (validated range: 0.1-2.0 Hz)
    freq_hz: float = Field(0.5, ge=0.1, le=2.0, description="Oscillation frequency")

    # Manual PHI0 override (degrees, 0-360)
    phi0_deg: Optional[float] = Field(None, ge=0.0, lt=360.0, description="Initial phase offset")

    # Multi-state control (cardinality + state letter)
    cardinality: Optional[Literal[2, 3, 4, 8]] = Field(None, description="Basin count (2/3/4/8)")
    state_letter: Optional[str] = Field(None, pattern="^[a-h]$", description="Helix64 state letter")

    # Authorization gate (required for PHI0 override or multi-state)
    control_consent: bool = Field(False, description="Explicit consent for basin control")

    # Legacy params (backward compat)
    hw_on: Optional[bool] = Field(False, description="Enable hardware measurement")

    @validator('state_letter')
    def validate_state_letter_cardinality(cls, v, values):
        """Ensure state_letter is valid for given cardinality."""
        if v is None:
            return v

        cardinality = values.get('cardinality')
        if cardinality is None:
            return v  # Will be set by auto-Δ

        # Map cardinality to valid letters
        valid_letters = {
            2: ['a', 'b'],
            3: ['a', 'b', 'c'],
            4: ['a', 'b', 'c', 'd'],
            8: ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'],
        }

        if v not in valid_letters[cardinality]:
            raise ValueError(f"state_letter '{v}' invalid for cardinality={cardinality}. Valid: {valid_letters[cardinality]}")

        return v


class CDCParams(BaseModel):
    """CDC regenerative model parameters."""

    # FibPi3D geometry
    fib_scale: Optional[float] = Field(None, gt=0.0, description="Fibonacci lattice scale")
    fib_shift_deg: Optional[float] = Field(None, description="Phase shift (degrees)")

    # Seed override (base64-encoded .hxz)
    seed_b64: Optional[str] = Field(None, description="Base64-encoded seed override")

    # Overlay schedule (Ξ(t) stepper)
    overlay_schedule: Optional[Dict[str, Any]] = Field(None, description="Phase modulation schedule")
    phase_omega: Optional[float] = Field(None, description="Angular frequency (rad/token)")

    # Timing params (for ω inference)
    freq_hz: Optional[float] = Field(None, ge=0.1, le=2.0, description="Frequency (Hz)")
    dt_ms_per_token: Optional[float] = Field(None, gt=0.0, description="Time per token (ms)")


class Message(BaseModel):
    """Chat message."""
    role: str = Field(..., pattern="^(user|assistant|system)$")
    content: str


class ChatRequest(BaseModel):
    """OpenAI-compatible chat completion request with Helix extensions."""

    # Standard OpenAI params
    model: Optional[str] = None
    messages: List[Message]
    max_tokens: int = Field(256, ge=1, le=4096)
    temperature: float = Field(0.7, ge=0.0, le=2.0)
    top_p: float = Field(0.9, ge=0.0, le=1.0)

    # Helix extensions
    quant: Optional[QuantOverlay] = None
    cdc: Optional[CDCParams] = None


class BasinMeasurement(BaseModel):
    """Phase-locked loop measurement result."""
    phi_deg: float = Field(..., ge=0.0, lt=360.0, description="Basin phase angle")
    R_spectral: float = Field(..., ge=0.0, le=1.0, description="Locking strength")
    MI_bits: Optional[float] = Field(None, ge=0.0, description="Mutual information")


class ControlMetadata(BaseModel):
    """Multi-state control provenance."""
    cardinality: int = Field(..., description="Basin count")
    state: str = Field(..., description="Target state letter")
    delta_base: float = Field(..., description="DELTA_BASE angle")
    phi0_applied: float = Field(..., description="Applied PHI0")


class EnergyMetrics(BaseModel):
    """Energy efficiency metrics."""
    uJ: float = Field(..., description="Total energy (microjoules)")
    uJ_per_nat: Optional[float] = Field(None, description="Energy per natural bit")


class SignedReceipt(BaseModel):
    """HMAC-signed proof of basin control."""

    timestamp: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    basin: BasinMeasurement
    control: ControlMetadata
    energy: EnergyMetrics
    signature: str = Field(..., description="HMAC-SHA256 signature")

    # Provenance
    schema_version: str = Field("helix:basin_control:v1", description="Schema version")
    git_commit: Optional[str] = Field(None, description="Git commit hash")


class HealthResponse(BaseModel):
    """Health endpoint response."""
    ok: bool
    engine_available: bool
    cdc_available: bool
    ga_nudge_forbid: bool
    uptime_seconds: float
    request_count: int
    error_count: int


class MetricsResponse(BaseModel):
    """Prometheus-compatible metrics."""
    flowtorch_requests_total: int
    flowtorch_errors_total: int
    flowtorch_basin_locks_total: int
    flowtorch_mi_sum: float
    flowtorch_latency_seconds_sum: float
    flowtorch_energy_uJ_sum: float


class CDSResponse(BaseModel):
    """Control Diagnostics Status - real-time steering health."""

    # MI degradation tracking
    recent_mi_values: List[float] = Field(description="Last N MI measurements")
    mean_mi: float = Field(description="Mean MI over window")
    mi_floor: float = Field(description="Minimum acceptable MI")
    mi_degraded: bool = Field(description="Below floor threshold")

    # Jitter budget
    recent_phi_jitter: List[float] = Field(description="Recent PHI0 jitter (degrees)")
    mean_jitter: float = Field(description="Mean jitter (degrees)")
    jitter_budget: float = Field(description="Max acceptable jitter (degrees)")
    jitter_exceeded: bool = Field(description="Jitter over budget")

    # Rate limiting
    requests_last_minute: int
    rate_limit: int
    rate_limited: bool

    # Overall status
    status: Literal["healthy", "degraded", "critical"]
