"""Tests for output schema validation.

Verifies that validate_experiment_output correctly checks
the canonical experiment output layout.
"""

import json
from pathlib import Path

import pytest

from clawbench.output_schema import validate_experiment_output


def _make_split_block(n_tasks: int = 2, score: float = 0.5) -> dict:
    """Build a minimal valid split block."""
    results = [
        {"task_name": f"task-{i}", "avg_reward": score, "passed": score > 0.5, "rewards": [score]}
        for i in range(n_tasks)
    ]
    passed = sum(1 for r in results if r["passed"])
    return {
        "avg_reward": score,
        "pass_rate": passed / n_tasks if n_tasks else 0,
        "summary": {"passed": passed, "total": n_tasks, "avg_reward": score},
        "results": results,
    }


def _write_valid_output(output_dir: Path, iterations: int = 2) -> None:
    """Write a minimal valid experiment output to output_dir."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # results.json
    results = {
        "engine": "custom",
        "iterations_run": iterations,
        "best_iteration": 1,
        "best_avg_reward": 0.7,
        "best_train_score": "1/2",
        "best_val_score": "1/2",
        "total_metric_calls": iterations * 4,
        "history": [
            {"iteration": i + 1, "train_score": "1/2", "val_score": "1/2",
             "train_avg_reward": 0.5, "val_avg_reward": 0.5}
            for i in range(iterations)
        ],
    }
    (output_dir / "results.json").write_text(json.dumps(results))

    # best_SKILL.md
    (output_dir / "best_SKILL.md").write_text("# Best skill")

    # skill_iter_0 (baseline)
    (output_dir / "skill_iter_0").mkdir()
    (output_dir / "skill_iter_0" / "SKILL.md").write_text("# Seed skill")

    # Per-iteration
    for n in range(1, iterations + 1):
        # iter_N_results.json
        iter_data = {"train": _make_split_block(), "val": _make_split_block()}
        (output_dir / f"iter_{n}_results.json").write_text(json.dumps(iter_data))

        # skill_iter_N/SKILL.md
        skill_dir = output_dir / f"skill_iter_{n}"
        skill_dir.mkdir()
        (skill_dir / "SKILL.md").write_text(f"# Skill iter {n}")


class TestValidateExperimentOutput:
    def test_valid_output(self, tmp_path):
        _write_valid_output(tmp_path)
        errors = validate_experiment_output(tmp_path)
        assert errors == []

    def test_missing_results_json(self, tmp_path):
        tmp_path.mkdir(exist_ok=True)
        errors = validate_experiment_output(tmp_path)
        assert any("results.json not found" in e for e in errors)

    def test_invalid_results_json(self, tmp_path):
        tmp_path.mkdir(exist_ok=True)
        (tmp_path / "results.json").write_text("not json")
        errors = validate_experiment_output(tmp_path)
        assert any("not valid JSON" in e for e in errors)

    def test_missing_required_keys(self, tmp_path):
        tmp_path.mkdir(exist_ok=True)
        (tmp_path / "results.json").write_text(json.dumps({"engine": "custom"}))
        (tmp_path / "best_SKILL.md").write_text("x")
        (tmp_path / "skill_iter_0").mkdir()
        (tmp_path / "skill_iter_0" / "SKILL.md").write_text("x")
        errors = validate_experiment_output(tmp_path)
        assert any("missing keys" in e for e in errors)

    def test_missing_baseline_skill(self, tmp_path):
        _write_valid_output(tmp_path)
        # Remove baseline
        (tmp_path / "skill_iter_0" / "SKILL.md").unlink()
        (tmp_path / "skill_iter_0").rmdir()
        errors = validate_experiment_output(tmp_path)
        assert any("skill_iter_0" in e for e in errors)

    def test_missing_best_skill(self, tmp_path):
        _write_valid_output(tmp_path)
        (tmp_path / "best_SKILL.md").unlink()
        errors = validate_experiment_output(tmp_path)
        assert any("best_SKILL.md" in e for e in errors)

    def test_missing_iter_results(self, tmp_path):
        _write_valid_output(tmp_path, iterations=2)
        (tmp_path / "iter_2_results.json").unlink()
        errors = validate_experiment_output(tmp_path)
        assert any("iter_2_results.json not found" in e for e in errors)

    def test_missing_iter_skill(self, tmp_path):
        _write_valid_output(tmp_path, iterations=2)
        (tmp_path / "skill_iter_2" / "SKILL.md").unlink()
        errors = validate_experiment_output(tmp_path)
        assert any("skill_iter_2" in e for e in errors)

    def test_iter_results_missing_split_keys(self, tmp_path):
        _write_valid_output(tmp_path, iterations=1)
        # Write iter results with incomplete split block
        (tmp_path / "iter_1_results.json").write_text(
            json.dumps({"train": {"avg_reward": 0.5}, "val": _make_split_block()})
        )
        errors = validate_experiment_output(tmp_path)
        assert any("train" in e and "missing keys" in e for e in errors)

    def test_valid_trajectories_accepted(self, tmp_path):
        _write_valid_output(tmp_path, iterations=1)
        traj = {
            "task-0": [{"score": 0.5, "steps": [], "summary": "did stuff", "run": 0, "duration_sec": 1.0}]
        }
        (tmp_path / "iter_1_trajectories.json").write_text(json.dumps(traj))
        errors = validate_experiment_output(tmp_path)
        assert errors == []

    def test_invalid_trajectory_entry(self, tmp_path):
        _write_valid_output(tmp_path, iterations=1)
        traj = {"task-0": [{"score": 0.5}]}  # missing steps, summary
        (tmp_path / "iter_1_trajectories.json").write_text(json.dumps(traj))
        errors = validate_experiment_output(tmp_path)
        assert any("missing keys" in e for e in errors)
