#!/usr/bin/env python3
"""Drive every enabled evaluation suite for one model, from one config file.

Suites that are disabled in the config are recorded as "Not run" with the reason
from the config, so an unexecuted benchmark is visible in the output rather than
absent from it.

    python evaluation/run_all.py --config evaluation/configs/piko_9b.yaml
    python evaluation/run_all.py --config evaluation/configs/base_model.yaml --dry-run
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]


def load_config(path: Path) -> dict[str, Any]:
    try:
        import yaml
    except ImportError:
        sys.exit("pyyaml is required: pip install pyyaml")
    if not path.is_file():
        sys.exit(f"Config not found: {path}")
    return yaml.safe_load(path.read_text(encoding="utf-8"))


def build_commands(config: dict[str, Any], results: Path) -> list[dict[str, Any]]:
    model = config["model"]
    runtime = config.get("runtime", {})
    generation = config.get("generation", {})
    label = model.get("label") or Path(model["id"]).name

    common = [
        "--model",
        model["id"],
        "--quantization",
        str(runtime.get("quantization", "4bit")),
        "--dtype",
        str(runtime.get("dtype", "bfloat16")),
    ]

    planned: list[dict[str, Any]] = []
    suites = config.get("suites", {})

    if suites.get("smoke", {}).get("enabled"):
        planned.append(
            {
                "suite": "smoke",
                "command": [
                    sys.executable,
                    str(REPO_ROOT / "evaluation" / "run_smoke_eval.py"),
                    "--config",
                    str(config["__path__"]),
                    "--output",
                    str(results / f"smoke_{label}.json"),
                ],
            }
        )

    custom = suites.get("custom_suite", {})
    if custom.get("enabled"):
        command = [
            sys.executable,
            str(REPO_ROOT / "evaluation" / "custom_suite" / "run_custom_eval.py"),
            *common,
            "--label",
            label,
            "--seed",
            str(generation.get("seed", 0)),
            "--max-new-tokens",
            str(generation.get("max_new_tokens", 384)),
            "--output",
            str(results / f"custom_suite_{label}.json"),
        ]
        if custom.get("categories") and custom["categories"] != "all":
            command += ["--category", str(custom["categories"])]
        planned.append({"suite": "custom_suite", "command": command})

    for name, block in suites.items():
        if name in ("smoke", "custom_suite"):
            continue
        if not block.get("enabled"):
            planned.append(
                {
                    "suite": name,
                    "command": None,
                    "status": "Not run",
                    "reason": block.get("note")
                    or "disabled in config (see evaluation/README.md for runtime cost)",
                    "limit": block.get("limit"),
                }
            )
    return planned


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--continue-on-error", action="store_true", default=True)
    args = parser.parse_args()

    config = load_config(args.config)
    config["__path__"] = str(args.config)
    label = config["model"].get("label") or Path(config["model"]["id"]).name
    results = REPO_ROOT / config.get("output", {}).get("directory", "evaluation/results")
    results.mkdir(parents=True, exist_ok=True)

    planned = build_commands(config, results)

    print(f"Model: {config['model']['id']}  (label: {label})")
    print(f"Results: {results}\n")
    for entry in planned:
        if entry["command"]:
            print(f"  RUN      {entry['suite']}")
        else:
            print(f"  NOT RUN  {entry['suite']}: {entry['reason']}")
    print()

    if args.dry_run:
        return

    manifest: list[dict[str, Any]] = []
    for entry in planned:
        if not entry["command"]:
            manifest.append(
                {
                    "suite": entry["suite"],
                    "status": "Not run",
                    "reason": entry["reason"],
                    "limit": entry.get("limit"),
                }
            )
            continue

        print(f"=== {entry['suite']} ===", flush=True)
        began = time.time()
        proc = subprocess.run(entry["command"], cwd=REPO_ROOT)
        seconds = round(time.time() - began, 1)

        record = {
            "suite": entry["suite"],
            "status": "completed" if proc.returncode == 0 else "failed",
            "returncode": proc.returncode,
            "seconds": seconds,
            "command": " ".join(entry["command"]),
        }
        manifest.append(record)
        print(f"--- {entry['suite']}: {record['status']} in {seconds}s\n", flush=True)

        if proc.returncode != 0 and not args.continue_on_error:
            break

    manifest_path = results / f"run_manifest_{label}.json"
    manifest_path.write_text(
        json.dumps(
            {
                "model": config["model"]["id"],
                "label": label,
                "config": str(args.config),
                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
                "suites": manifest,
            },
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    print(f"wrote {manifest_path}")

    failed = [m for m in manifest if m.get("status") == "failed"]
    if failed:
        print(f"\n{len(failed)} suite(s) failed: {[m['suite'] for m in failed]}")
        sys.exit(1)


if __name__ == "__main__":
    main()