"""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

        if str(path).endswith('.jsonl'):
            # streaming format: first line = report header, then one
            # sample per line (grep/split/tail friendly)
            head = self.model_dump(exclude={'samples'})
            head['type'] = 'report'
            with open(path, 'w', encoding='utf-8') as f:
                f.write(json.dumps(head, ensure_ascii=False) + '\n')
                for smp in self.samples:
                    row = smp if isinstance(smp, dict) else smp.model_dump()
                    row['type'] = 'sample'
                    f.write(json.dumps(row, ensure_ascii=False) + '\n')
            return
        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

        if str(path).endswith('.jsonl'):
            head, samples = None, []
            with open(path, encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    row = json.loads(line)
                    if row.pop('type', '') == 'report' or head is None:
                        head = {k: v for k, v in row.items() if k != 'type'}
                    else:
                        samples.append({k: v for k, v in row.items() if k != 'type'})
            head = head or {}
            head['samples'] = samples
            return cls.model_validate(head)
        with open(path, encoding='utf-8') as f:
            return cls.model_validate(json.load(f))