"""Canonical experiment output schema and validation.

Both the custom engine (custom_loop.py) and GEPA engine (optimize.py)
must produce output conforming to this schema so the dashboard API and UI
work identically regardless of engine.
"""

from __future__ import annotations

import json
from pathlib import Path

# ---------------------------------------------------------------------------
# Schema definitions (required keys)
# ---------------------------------------------------------------------------

RESULTS_JSON_REQUIRED_KEYS = {
    "engine",
    "iterations_run",
    "best_iteration",
    "best_avg_reward",
    "history",
}

ITER_RESULTS_JSON_REQUIRED_KEYS = {
    "train",
    "val",
}

SPLIT_BLOCK_REQUIRED_KEYS = {
    "avg_reward",
    "pass_rate",
    "summary",
    "results",
}

ITER_TRAJECTORIES_ENTRY_REQUIRED_KEYS = {
    "score",
    "steps",
    "summary",
}


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


class OutputValidationError(Exception):
    """Raised when experiment output does not conform to the schema."""

    def __init__(self, errors: list[str]):
        self.errors = errors
        super().__init__(f"Output validation failed with {len(errors)} error(s):\n" +
                         "\n".join(f"  - {e}" for e in errors))


def validate_experiment_output(output_dir: str | Path) -> list[str]:
    """Validate that an experiment output directory conforms to the schema.

    Checks:
    - results.json exists and has required keys
    - For each iteration 1..N: iter_N_results.json exists with required keys,
      skill_iter_N/SKILL.md exists
    - skill_iter_0/SKILL.md (baseline) exists
    - best_SKILL.md exists

    Returns:
        List of error strings (empty means valid).
    """
    output_dir = Path(output_dir)
    errors: list[str] = []

    # 1. results.json
    results_path = output_dir / "results.json"
    if not results_path.exists():
        errors.append("results.json not found")
        return errors  # can't continue without it

    try:
        results = json.loads(results_path.read_text())
    except json.JSONDecodeError as e:
        errors.append(f"results.json is not valid JSON: {e}")
        return errors

    missing = RESULTS_JSON_REQUIRED_KEYS - set(results.keys())
    if missing:
        errors.append(f"results.json missing keys: {sorted(missing)}")

    iterations_run = results.get("iterations_run", 0)

    # 2. best_SKILL.md
    if not (output_dir / "best_SKILL.md").exists():
        errors.append("best_SKILL.md not found")

    # 3. Baseline skill (iter 0)
    if not (output_dir / "skill_iter_0" / "SKILL.md").exists():
        errors.append("skill_iter_0/SKILL.md (baseline) not found")

    # 4. Per-iteration files
    for n in range(1, iterations_run + 1):
        # iter_N_results.json
        iter_results_path = output_dir / f"iter_{n}_results.json"
        if not iter_results_path.exists():
            errors.append(f"iter_{n}_results.json not found")
        else:
            try:
                iter_data = json.loads(iter_results_path.read_text())
                iter_missing = ITER_RESULTS_JSON_REQUIRED_KEYS - set(iter_data.keys())
                if iter_missing:
                    errors.append(f"iter_{n}_results.json missing keys: {sorted(iter_missing)}")
                else:
                    # Check split blocks
                    for split in ("train", "val"):
                        block = iter_data.get(split, {})
                        block_missing = SPLIT_BLOCK_REQUIRED_KEYS - set(block.keys())
                        if block_missing:
                            errors.append(
                                f"iter_{n}_results.json['{split}'] missing keys: {sorted(block_missing)}"
                            )
            except json.JSONDecodeError as e:
                errors.append(f"iter_{n}_results.json is not valid JSON: {e}")

        # skill_iter_N/SKILL.md
        if not (output_dir / f"skill_iter_{n}" / "SKILL.md").exists():
            errors.append(f"skill_iter_{n}/SKILL.md not found")

        # iter_N_trajectories.json is optional — but if present, validate structure
        traj_path = output_dir / f"iter_{n}_trajectories.json"
        if traj_path.exists():
            try:
                traj_data = json.loads(traj_path.read_text())
                if not isinstance(traj_data, dict):
                    errors.append(f"iter_{n}_trajectories.json must be a dict (task_name -> entries)")
                else:
                    for task_name, entries in traj_data.items():
                        if not isinstance(entries, list):
                            errors.append(
                                f"iter_{n}_trajectories.json['{task_name}'] must be a list"
                            )
                            continue
                        for j, entry in enumerate(entries):
                            entry_missing = ITER_TRAJECTORIES_ENTRY_REQUIRED_KEYS - set(entry.keys())
                            if entry_missing:
                                errors.append(
                                    f"iter_{n}_trajectories.json['{task_name}'][{j}] "
                                    f"missing keys: {sorted(entry_missing)}"
                                )
            except json.JSONDecodeError as e:
                errors.append(f"iter_{n}_trajectories.json is not valid JSON: {e}")

    return errors
