from __future__ import annotations

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field
from typing import Literal


class PredictionReceipt(BaseModel):
    """
    Horizon prediction snapshot.
    
    Latent forecasts at fast/mid/slow timescales. Latents themselves should be stored
    as references or hashes (avoid embedding large blobs in receipts).
    """

    receipt_type: Literal["PredictionReceipt"] = "PredictionReceipt"

    receipt_id: str
    ts: Any  # datetime or ISO string
    envelope_id: Optional[str] = None
    envelope_hash: Optional[str] = None
    parent_ids: List[str] = Field(default_factory=list)
    domain: Optional[str] = None
    profile: Optional[str] = None
    risk_band: Optional[str] = None

    state_id: str
    t_start: Optional[float] = None
    horizon: int
    fast_refs: List[str] = Field(default_factory=list)
    mid_refs: List[str] = Field(default_factory=list)
    slow_refs: List[str] = Field(default_factory=list)
    slots_summary: Dict[str, Any] = Field(default_factory=dict)
    action_seq_id: Optional[str] = None
    value_estimate: Optional[float] = None
    status: str = Field(default="ok")  # ok|partial|error
    error_message: Optional[str] = None


class SurpriseReceipt(BaseModel):
    """
    Horizon prediction error snapshot (surprise).
    
    Captures deviation between predicted and observed latent states. Thresholds and
    alert flag are included for downstream gating/telemetry.
    """

    receipt_type: Literal["SurpriseReceipt"] = "SurpriseReceipt"

    receipt_id: str
    ts: Any  # datetime or ISO string
    envelope_id: Optional[str] = None
    envelope_hash: Optional[str] = None
    parent_ids: List[str] = Field(default_factory=list)
    domain: Optional[str] = None
    profile: Optional[str] = None
    risk_band: Optional[str] = None

    state_id: str
    prediction_ref: Optional[str] = None
    t_observed: Optional[float] = None
    error_fast: Optional[float] = None
    error_mid: Optional[float] = None
    error_slow: Optional[float] = None
    slot_error: Dict[str, Any] = Field(default_factory=dict)
    channels: Dict[str, Any] = Field(default_factory=dict)
    status: str = Field(default="ok")  # ok|warning|error
    thresholds: Dict[str, float] = Field(default_factory=dict)  # e.g., {"warn":0.5,"alert":0.8}
    alert: bool = False
    error_message: Optional[str] = None


class WorldHypothesisReceipt(BaseModel):
    """
    Horizon counterfactual/hypothesis snapshot.
    
    Represents imagined rollout endpoints for a candidate action sequence. Endpoint
    latents should be referenced by id/hash, not embedded.
    """

    receipt_type: Literal["WorldHypothesisReceipt"] = "WorldHypothesisReceipt"

    receipt_id: str
    ts: Any  # datetime or ISO string
    envelope_id: Optional[str] = None
    envelope_hash: Optional[str] = None
    parent_ids: List[str] = Field(default_factory=list)
    domain: Optional[str] = None
    profile: Optional[str] = None
    risk_band: Optional[str] = None

    state_id: str
    action_seq_id: str
    horizon: int
    value_estimate: Optional[float] = None
    safety_estimate: Optional[float] = None
    fast_endpoint_ref: Optional[str] = None
    mid_endpoint_ref: Optional[str] = None
    slow_endpoint_ref: Optional[str] = None
    status: str = Field(default="ok")  # ok|partial|error
    error_message: Optional[str] = None
