"""Build the distillation baseline the reality check compares reinforcement learning against.

Purpose: generate solutions to the same training problems with a larger local teacher,
    keep only the ones the verifier says are correct, and write them in the shape Part 11's
    train-sft.py reads. The point is a fair comparison: the reality check asks whether
    GRPO's gain is worth its budget, and the honest way to answer that is to spend a
    similar budget on the simplest alternative and measure both. Part 15 teaches
    distillation properly; this is the one-page version of it, used as a control.
Platform: all (pure Python over HTTP; the teacher runs on whatever engine your track uses)
Minimum memory: 16 GB on the machine serving the teacher; this script needs very little
Assumes: Python 3.10 or newer; an OpenAI-compatible endpoint serving a larger model than
    the student, such as llama-server from Part 6 or the gateway from Part 9;
    make-tasks.py has written the training problems; rewards.py and runlog.py sit next
    to this file.

Usage: python3 distil-baseline.py --tasks tasks/train.jsonl --base-url http://127.0.0.1:8080/v1 \
           --teacher-model qwen3-8b --out-dir distil-data --labbook labbook.md
       python3 distil-baseline.py --tasks tasks/train.jsonl --base-url http://127.0.0.1:8080/v1 \
           --teacher-model qwen3-8b --attempts 2 --budget-seconds 1800 --out-dir distil-data
"""

from __future__ import annotations

import argparse
import hashlib
import json
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Optional

import rewards as reward_lib
import runlog


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def write_jsonl(path: Path, rows: list[dict]) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--tasks", default="tasks/train.jsonl", help="the same problems GRPO trained on")
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--teacher-model", required=True, help="a larger model than the student")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--attempts", type=int, default=1,
                        help="samples per problem; extra attempts recover problems the teacher got wrong once")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--budget-seconds", type=float, default=None,
                        help="stop generating after this long, so the comparison is at equal wall clock")
    parser.add_argument("--valid-fraction", type=float, default=0.1)
    parser.add_argument("--tolerance", type=float, default=1e-6)
    parser.add_argument("--keep-wrong", action="store_true",
                        help="keep solutions the verifier rejects; off by default, and the lesson says why")
    parser.add_argument("--out-dir", default="distil-data")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    rows = [json.loads(line) for line in Path(args.tasks).read_text(encoding="utf-8").splitlines() if line.strip()]
    print(f"{len(rows)} problems; teacher {args.teacher_model} at {args.base_url}")

    endpoint = args.base_url.rstrip("/") + "/chat/completions"
    grade = reward_lib.numeric_reward(args.tolerance, fallback_to_last_number=False)

    kept: list[dict] = []
    attempted = 0
    correct = 0
    teacher_completion_tokens = 0
    started = time.time()
    stopped_early = False

    for i, row in enumerate(rows, start=1):
        if args.budget_seconds and time.time() - started > args.budget_seconds:
            stopped_early = True
            print(f"budget of {args.budget_seconds:.0f} s reached after {i - 1} problems")
            break
        for attempt in range(args.attempts):
            payload = {
                "model": args.teacher_model,
                "messages": row["prompt"],
                "temperature": args.temperature,
                "max_tokens": args.max_tokens,
                "seed": args.seed + attempt,
            }
            body = post_json(endpoint, payload, args.api_key, args.timeout)
            text = (body["choices"][0]["message"]["content"] or "").strip()
            teacher_completion_tokens += int(body.get("usage", {}).get("completion_tokens") or 0)
            attempted += 1
            is_right = grade(completions=[text], answer=[row["answer"]])[0] == 1.0
            if is_right:
                correct += 1
            if is_right or args.keep_wrong:
                kept.append({
                    "prompt": row["prompt"],
                    "completion": [{"role": "assistant", "content": text}],
                })
                break
        print(f"  [{i}/{len(rows)}] kept {len(kept)}  teacher correct {correct}/{attempted}")

    if not kept:
        raise SystemExit("the teacher solved none of the problems in the required format. Check that "
                         "it is a larger model than the student and that the endpoint is right.")

    split = max(1, int(len(kept) * args.valid_fraction))
    valid, train = kept[:split], kept[split:]
    out = Path(args.out_dir)
    train_hash = write_jsonl(out / "train.jsonl", train)
    write_jsonl(out / "valid.jsonl", valid)

    elapsed = time.time() - started
    budget = {
        "teacher_model": args.teacher_model,
        "problems_attempted": attempted,
        "teacher_correct": correct,
        "teacher_accuracy": round(correct / attempted, 4) if attempted else None,
        "examples_kept": len(kept),
        "teacher_completion_tokens": teacher_completion_tokens or None,
        "generation_seconds": round(elapsed, 1),
        "stopped_on_budget": stopped_early,
    }
    print("\n" + json.dumps(budget, indent=2))
    print(f"\nwritten to {out}/  train {len(train)}  valid {len(valid)}  sha256 {train_hash[:12]}...")
    print("Now train the same student on it with the Part 11 recipe, unchanged:")
    print(f"  python3 train-sft.py --model <the same base model> --data-dir {out} \\")
    print("      --output-dir runs/distil-baseline --epochs 3 --labbook labbook.md")
    print("Then score it with eval-pass-at-1.py on the same held-out files as the GRPO run.")
    if teacher_completion_tokens == 0:
        print("\nThe server returned no token usage, so the budget line records wall clock only.")

    if args.labbook:
        record = runlog.record(
            labbook=args.labbook,
            lab="part-14/distil-baseline",
            model=args.teacher_model,
            dataset={"path": args.tasks, "sha256": runlog.file_sha256(args.tasks),
                     "problems": len(rows), "examples_kept": len(kept),
                     "written_sha256": train_hash},
            hyperparameters={"attempts": args.attempts, "temperature": args.temperature,
                             "max_tokens": args.max_tokens, "keep_wrong": args.keep_wrong,
                             "budget_seconds": args.budget_seconds},
            seed=args.seed,
            losses={},
            scores=budget,
            config_path=__file__,
            notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()