"""Run the whole distillation pipeline from one configuration file, and record every stage.

Purpose: the project's runner. Reads distil-config.yaml, turns each stage into the
    command it stands for, runs the stages you asked for, and writes one pipeline
    record per stage into the lab notebook carrying the hash of the configuration
    that produced it. Stages are resumable and skippable, so a run interrupted in
    the third hour continues rather than restarting, and a changed threshold reruns
    the filter without regenerating a dataset.
Platform: all (pure Python; the stages it launches have their own platform notes,
    and Track M's training stage prints an mlx_lm.lora command instead of running
    the PyTorch one)
Minimum memory: 12 GB, set by the stages rather than by this script
Assumes: Python 3.10 or newer, PyYAML installed, and this part's scripts in the same
    directory. The teacher is already served; this script starts no engine, because
    a runner that starts models is a runner that hides which model answered.

Usage: python3 pipeline.py --config distil-config.yaml --dry-run
       python3 pipeline.py --config distil-config.yaml
       python3 pipeline.py --config distil-config.yaml --from train
       python3 pipeline.py --config distil-config.yaml --only evaluate
       python3 pipeline.py --config distil-config.yaml --run-export
       python3 pipeline.py --config distil-config.yaml --status

--dry-run prints every command and touches nothing, which is the way to read a
configuration change before spending an evening on it.
"""

from __future__ import annotations

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

import distillog

ORDER = ["seed_prompts", "generate", "filter", "train", "evaluate", "export"]


def load_config(path: Path) -> dict:
    try:
        import yaml  # noqa: PLC0415 - reported clearly rather than crashing on import
    except ImportError:
        raise SystemExit(
            "PyYAML is not installed in this environment. Install it with "
            "`uv pip install pyyaml`, or read the configuration by hand and run the "
            "stage scripts directly; every stage is one command."
        ) from None
    if not path.is_file():
        raise SystemExit(f"{path} not found")
    config = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(config, dict):
        raise SystemExit(f"{path} did not parse as a mapping")
    for required in ("run", "teacher", "student", "stages"):
        if required not in config:
            raise SystemExit(f"{path} has no '{required}' section")
    return config


def expand(value: Any) -> Any:
    """Expand ~ in anything that looks like a path, leaving everything else alone."""
    if isinstance(value, str) and value.startswith("~"):
        return os.path.expanduser(value)
    return value


def flag(name: str, value: Any) -> list[str]:
    """One configuration key as command-line arguments, or nothing when it is unset."""
    if value is None or value == "":
        return []
    option = "--" + name.replace("_", "-")
    if isinstance(value, bool):
        return [option] if value else []
    return [option, str(expand(value))]


def build_commands(config: dict, here: Path) -> dict[str, list[list[str]]]:
    """Every stage as the list of commands it stands for. Nothing is run here."""
    run = config["run"]
    teacher = config["teacher"]
    student = config["student"]
    labbook = run.get("labbook", "labbook.md")
    py = sys.executable
    commands: dict[str, list[list[str]]] = {}

    seeds = config.get("seed_prompts", {})
    if seeds:
        commands["seed_prompts"] = [[
            py, str(here / seeds.get("script", "make-seed-prompts.py")),
            *flag("out_dir", seeds.get("out_dir", "seeds")),
            *flag("train", seeds.get("train")),
            *flag("heldout", seeds.get("heldout")),
            *flag("maths", seeds.get("maths")),
            *flag("seed", run.get("seed", 0)),
        ]]

    gen = config.get("generate", {})
    if gen:
        commands["generate"] = [[
            py, str(here / gen.get("script", "generate-teacher-data.py")),
            *flag("seeds", gen.get("seeds")),
            *flag("out", gen.get("out")),
            *flag("base_url", teacher.get("base_url")),
            *flag("model", teacher.get("served_as")),
            *flag("teacher_id", teacher.get("id")),
            *flag("quant", teacher.get("quant")),
            *flag("thinking", gen.get("thinking", False)),
            *flag("temperature", gen.get("temperature")),
            *flag("top_p", gen.get("top_p")),
            *flag("max_tokens", gen.get("max_tokens")),
            *flag("samples", gen.get("samples")),
            *flag("concurrency", gen.get("concurrency")),
            *flag("resume", gen.get("resume", True)),
            *flag("labbook", labbook),
        ]]

    filt = config.get("filter", {})
    if filt:
        commands["filter"] = [[
            py, str(here / filt.get("script", "filter-and-dedupe.py")),
            *flag("raw", filt.get("raw")),
            *flag("out_dir", filt.get("out_dir", ".")),
            *flag("tasks", filt.get("tasks")),
            *flag("min_words", filt.get("min_words")),
            *flag("max_words", filt.get("max_words")),
            *flag("near_duplicate", filt.get("near_duplicate")),
            *flag("contamination", filt.get("contamination")),
            *flag("keep_thinking", filt.get("keep_thinking", False)),
            *flag("valid_fraction", filt.get("valid_fraction")),
            *flag("seed", run.get("seed", 0)),
            *flag("labbook", labbook),
        ]]

    train = config.get("train", {})
    if train:
        if student.get("trainer") == "mlx":
            # Track M: print the documented mlx-lm command rather than running the
            # PyTorch path, which would silently train a different thing.
            commands["train"] = [[
                "mlx_lm.lora",
                "--model", str(student.get("hf_id")),
                "--train",
                "--data", str(train.get("data_dir", "data")) + "-mlx",
                "--batch-size", str(train.get("batch_size", 2)),
                "--iters", str(train.get("iters", 600)),
                "--adapter-path", str(train.get("output_dir", "runs/seq-student")),
            ]]
        else:
            commands["train"] = [[
                py, str(here / train.get("script", "train-student.py")),
                *flag("model", student.get("hf_id")),
                *flag("teacher_id", teacher.get("id")),
                *flag("data_dir", train.get("data_dir")),
                *flag("output_dir", train.get("output_dir")),
                *flag("epochs", train.get("epochs")),
                *flag("batch_size", train.get("batch_size")),
                *flag("grad_accum", train.get("grad_accum")),
                *flag("lr", train.get("learning_rate")),
                *flag("max_length", train.get("max_length")),
                *flag("rank", train.get("rank")),
                *flag("alpha", train.get("alpha")),
                *flag("gradient_checkpointing", train.get("gradient_checkpointing", False)),
                *flag("seed", run.get("seed", 0)),
                *flag("labbook", labbook),
            ]]

    ev = config.get("evaluate", {})
    if ev:
        commands["evaluate"] = [[
            py, str(here / ev.get("script", "evaluate-triplet.py")),
            *flag("harness_dir", ev.get("harness_dir")),
            *flag("tasks", ev.get("tasks")),
            *flag("base_url", teacher.get("base_url")),
            *flag("teacher", ev.get("teacher_served_as")),
            *flag("base_student", ev.get("base_student_served_as")),
            *flag("distilled", ev.get("distilled_served_as")),
            *flag("teacher_id", teacher.get("id")),
            *flag("student_id", student.get("id")),
            *flag("quant", ev.get("quant")),
            *flag("engine", ev.get("engine")),
            *flag("engine_version", ev.get("engine_version")),
            *flag("judge_model", ev.get("judge_model")),
            *flag("out_dir", ev.get("out_dir")),
            *flag("no_pause", ev.get("no_pause", False)),
            *flag("labbook", labbook),
        ]]

    exp = config.get("export", {})
    if exp:
        merge = expand(exp.get("merge_script", "merge-adapter.py"))
        adapter = (config.get("train", {}) or {}).get("output_dir", "runs/seq-student")
        out_dir = expand(exp.get("out_dir", "models/student"))
        llama_cpp = expand(exp.get("llama_cpp", "~/llama.cpp"))
        quant = exp.get("quant", "Q4_K_M")
        commands["export"] = [
            # Part 13's merge-adapter.py takes --adapter and --merged-dir; the
            # converter and the quantiser are llama.cpp's, as in Part 11's lab.
            [py, str(merge), "--adapter", str(adapter), "--merged-dir", str(out_dir)],
            [py, f"{llama_cpp}/convert_hf_to_gguf.py", str(out_dir),
             "--outfile", f"{out_dir}.gguf", "--outtype", "bf16"],
            [f"{llama_cpp}/build/bin/llama-quantize", f"{out_dir}.gguf",
             f"{out_dir}-{quant}.gguf", quant],
        ]

    return commands


def run_command(command: list[str], dry: bool) -> float:
    print("+ " + " ".join(command))
    if dry:
        return 0.0
    if shutil.which(command[0]) is None and not Path(command[0]).exists():
        raise SystemExit(f"{command[0]} not found on this machine")
    started = time.time()
    subprocess.run(command, check=True)
    return time.time() - started


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--config", default="distil-config.yaml")
    parser.add_argument("--dry-run", action="store_true", help="print the commands and change nothing")
    parser.add_argument("--from", dest="from_stage", default=None, help="start at this stage")
    parser.add_argument("--only", default=None, help="run exactly this stage")
    parser.add_argument("--run-export", action="store_true",
                        help="actually run the export commands rather than printing them")
    parser.add_argument("--status", action="store_true",
                        help="list the stages already recorded in the lab notebook and exit")
    args = parser.parse_args()

    here = Path(__file__).resolve().parent
    config_path = Path(args.config)
    config = load_config(config_path)
    labbook = config["run"].get("labbook", "labbook.md")

    # The hash is of the parsed configuration, not of the file's bytes, so a comment
    # or a reordering does not look like a different experiment while a changed
    # threshold does.
    config_hash = distillog.config_sha256(config)
    print(f"configuration: {config_path}  hash {config_hash[:16]}")
    print(f"run: {config['run'].get('name')} v{config['run'].get('version')}  "
          f"seed {config['run'].get('seed')}")

    if args.status:
        recorded = distillog.read_stages(labbook)
        if not recorded:
            print(f"no Part 15 stages recorded in {labbook} yet")
            return
        print(f"\n{'stage':<12}{'date':<22}{'run id':<26}notes")
        for record in recorded:
            print(f"{record.get('stage', '?'):<12}{record.get('date', '?'):<22}"
                  f"{record.get('run_id', '?'):<26}{record.get('notes') or ''}")
        return

    commands = build_commands(config, here)
    wanted = [s for s in ORDER if s in commands]
    configured = config.get("stages") or []
    wanted = [s for s in wanted if s in configured or s == "seed_prompts"]
    if args.only:
        if args.only not in commands:
            raise SystemExit(f"no stage named {args.only!r}; known: {', '.join(commands)}")
        wanted = [args.only]
    elif args.from_stage:
        if args.from_stage not in wanted:
            raise SystemExit(f"cannot start at {args.from_stage!r}; stages to run: {', '.join(wanted)}")
        wanted = wanted[wanted.index(args.from_stage):]

    print(f"stages: {', '.join(wanted)}")
    print()

    for stage in wanted:
        print(f"--- {stage} " + "-" * (60 - len(stage)))
        dry = args.dry_run or (stage == "export" and not args.run_export)
        if stage == "export" and not args.run_export and not args.dry_run:
            print("(printing only: pass --run-export to run these; they write large files)")
        elapsed = 0.0
        for command in commands[stage]:
            elapsed += run_command(command, dry)
        if dry:
            print()
            continue

        # The stage scripts write their own detailed records. This one is the
        # pipeline's own: which stage ran, under which configuration, for how long.
        distillog.record(
            labbook=labbook,
            lab="part-15/pipeline",
            stage=f"pipeline:{stage}",
            teacher={"id": config["teacher"].get("id"),
                     "served_as": config["teacher"].get("served_as")},
            student={"id": config["student"].get("id"),
                     "trainer": config["student"].get("trainer")},
            dataset={"config": str(config_path), "config_sha256": config_hash},
            hyperparameters={"run_name": config["run"].get("name"),
                             "run_version": config["run"].get("version"),
                             "commands": [" ".join(c) for c in commands[stage]]},
            seed=config["run"].get("seed"),
            cost=distillog.build_cost(seconds=elapsed),
            config_path=str(config_path),
            notes=None,
        )
        print(f"({stage} took {elapsed / 60:.1f} minutes)\n")

    print("done.")
    print(f"Run `python3 pipeline.py --config {config_path} --status` to see every stage "
          "this configuration has produced, and `python3 distillog.py --stages` for the "
          "detail each stage recorded.")
    print(json.dumps({"config_sha256": config_hash,
                      "stages_run": wanted,
                      "dry_run": args.dry_run}, indent=2))


if __name__ == "__main__":
    main()