"""Evaluation-layer record schemas.

The eval layer produces immutable artifacts: for every sample we keep the
RAW prediction (never the extracted string alone) so any recipe change can
re-score without re-running the model. Aggregated reports carry the recipe
fingerprint for reproducibility.

Data flow:  Dataset x predictions -> SampleResult per sample
            -> SampleResult list -> EvalReport (aggregates + artifacts)
"""

import time
from typing import Any, Dict, List, Optional, Union

from pydantic import BaseModel, Field


class SampleResult(BaseModel):
    """One evaluated sample. raw_prediction is the source of truth."""

    sample_id: Optional[int] = None
    dataset: str = ''
    subset: str = ''
    task_type: Optional[str] = None

    raw_prediction: str = ''                  # exactly what the model produced
    extracted_prediction: str = ''            # extractor output (derivable)
    extraction_ok: bool = True                # False => extractor found nothing usable
    extraction_note: str = ''                 # e.g. which cascade stage hit

    scores: Dict[str, float] = Field(default_factory=dict)      # {'acc': 1.0, 'f1': 0.6}
    score_details: Dict[str, Any] = Field(default_factory=dict) # {'acc': {'judge_raw': 'GRADE: C'}}

    target: Union[str, List[str]] = ''
    group_key: str = ''                       # pass@k task id / binned bucket / category
    metadata: Dict[str, Any] = Field(default_factory=dict)
    error: str = ''                           # scorer/extractor exception (never silently dropped)

    # agent hinge: when a future AgentLoop produces multi-turn trajectories,
    # they land here (list of ChatMessage dumps) + final environment state;
    # env_reward scorers consume these instead of extracted text
    trajectory: Optional[List[Dict[str, Any]]] = None
    env_state: Optional[Dict[str, Any]] = None
    usage: Optional[Dict[str, Any]] = None    # tokens/cost/latency per sample


class EvalReport(BaseModel):
    """Aggregated artifact: what a viewer/visualizer consumes."""

    dataset: str
    recipe: str = ''                          # e.g. 'gsm8k'
    recipe_version: str = ''
    model: str = ''                           # model name/url tag
    created_at: str = Field(default_factory=lambda: time.strftime('%Y-%m-%d %H:%M:%S'))

    num_samples: int = 0
    num_failed_extractions: int = 0
    metrics: Dict[str, float] = Field(default_factory=dict)     # {'acc': 0.62}
    metric_groups: Dict[str, Dict[str, float]] = Field(default_factory=dict)
    # {'by_category': {'algebra': 0.7, ...}, 'pass_at_k': {'pass@1': .., 'pass@8': ..},
    #  'by_length_bin': {'8k': .., '32k': ..}}

    samples: List[SampleResult] = Field(default_factory=list)   # full per-sample detail

    def save(self, path) -> None:
        import json

        with open(path, 'w', encoding='utf-8') as f:
            json.dump(self.model_dump(), f, ensure_ascii=False, indent=2)

    @classmethod
    def load(cls, path) -> 'EvalReport':
        import json

        with open(path, encoding='utf-8') as f:
            return cls.model_validate(json.load(f))