"""Answer a seed prompt file with your local teacher, concurrently, and resumably.

Purpose: the generation stage of sequence-level distillation. Sends every prompt in
    a seed file to an OpenAI-compatible endpoint with stated sampling settings and a
    fixed number of concurrent requests, writes one JSON line per answer as it
    arrives, and records the teacher id, the settings, the token counts and the
    wall-clock time. Re-running it skips prompts already answered, so a run
    interrupted after two hours resumes rather than restarting.
Platform: all (pure Python over HTTP; the teacher may be llama-server from Part 6,
    vLLM from Part 9, the Part 9 gateway, or an mlx-lm server on Track M)
Minimum memory: 12 GB on the machine serving the teacher; this script needs very little
Assumes: Python 3.10 or newer and a reachable OpenAI-compatible endpoint. The seed
    file is JSON Lines with an "id" and a "prompt" on every line, as written by
    make-seed-prompts.py. distillog.py sits next to this file.

Usage: python3 generate-teacher-data.py --seeds seeds/prompts.jsonl \\
           --base-url http://127.0.0.1:4000/v1 --model local/chat \\
           --teacher-id qwen3-30b-a3b --out raw/teacher.jsonl --labbook labbook.md
       python3 generate-teacher-data.py --seeds seeds/prompts.jsonl --out raw/teacher.jsonl \\
           --concurrency 8 --temperature 0.7 --top-p 0.8 --samples 2 --resume
       python3 generate-teacher-data.py --seeds seeds/maths.jsonl --out raw/maths-traces.jsonl \\
           --thinking --temperature 0.6 --top-p 0.95 --samples 4

The settings matter and are recorded on every line. Qwen3's model card gives
Temperature 0.6, TopP 0.95, TopK 20, MinP 0 for thinking mode and Temperature 0.7,
TopP 0.8, TopK 20, MinP 0 for non-thinking mode, and says not to use greedy decoding
in thinking mode; --thinking switches this script's defaults to the first set.
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Optional

import distillog

# Sampling defaults, from the Qwen3 model cards read on 2026-09-09. They are
# arguments rather than constants because a different teacher wants different
# numbers, and because a run whose settings are not recorded is not reproducible.
NON_THINKING = {"temperature": 0.7, "top_p": 0.8}
THINKING = {"temperature": 0.6, "top_p": 0.95}


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 read_jsonl(path: Path) -> list[dict]:
    rows = []
    with path.open(encoding="utf-8") as handle:
        for number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{number}: not valid JSON ({exc})") from exc
    return rows


def already_done(path: Path) -> set[str]:
    """Keys of the (prompt, sample) pairs already in the output file.

    Resumability is the difference between a generation stage you can run on a
    laptop overnight and one you have to babysit. The key includes the sample
    index so that --samples 4 resumes at the right sample, not the right prompt.
    """
    if not path.is_file():
        return set()
    done = set()
    for row in read_jsonl(path):
        if "id" in row and "sample" in row:
            done.add(f"{row['id']}#{row['sample']}")
    return done


class PowerSampler:
    """Samples accelerator power in the background so the run can report energy.

    Track M reports nothing to an unprivileged process, so mean() returns None
    there and the cost block records a null rather than a zero.
    """

    def __init__(self, interval: float = 15.0) -> None:
        self.interval = interval
        self.samples: list[float] = []
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None

    def start(self) -> None:
        if distillog.sample_power() is None:
            return
        self._thread = threading.Thread(target=self._loop, daemon=True)
        self._thread.start()

    def _loop(self) -> None:
        while not self._stop.wait(self.interval):
            value = distillog.sample_power()
            if value is not None:
                self.samples.append(value)

    def stop(self) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=2)

    def mean(self) -> float | None:
        return round(statistics.fmean(self.samples), 1) if self.samples else None


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--seeds", required=True, help="JSON Lines file of prompts")
    parser.add_argument("--out", required=True, help="JSON Lines file of teacher answers")
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--model", default="local/chat",
                        help="the name the endpoint answers to, e.g. a gateway alias")
    parser.add_argument("--teacher-id", default=None,
                        help="the course model id of the teacher, recorded on every line. "
                             "The API cannot be asked which weights are loaded, so say so here.")
    parser.add_argument("--quant", default="unknown",
                        help="the teacher's quantisation, recorded on every line")
    parser.add_argument("--system", default=None, help="system prompt sent with every request")
    parser.add_argument("--samples", type=int, default=1,
                        help="answers per prompt; more than one feeds rejection sampling")
    parser.add_argument("--thinking", action="store_true",
                        help="use the model card's thinking-mode sampling settings")
    parser.add_argument("--temperature", type=float, default=None)
    parser.add_argument("--top-p", type=float, default=None)
    parser.add_argument("--top-k", type=int, default=20)
    parser.add_argument("--max-tokens", type=int, default=768)
    parser.add_argument("--concurrency", type=int, default=4,
                        help="requests in flight. Match it to the server's slot count: "
                             "llama-server's --parallel, or vLLM's scheduler.")
    parser.add_argument("--limit", type=int, default=None, help="stop after this many prompts")
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--resume", action="store_true",
                        help="skip prompts already present in --out and append")
    parser.add_argument("--labbook", default=None)
    args = parser.parse_args()

    defaults = THINKING if args.thinking else NON_THINKING
    settings = {
        "temperature": args.temperature if args.temperature is not None else defaults["temperature"],
        "top_p": args.top_p if args.top_p is not None else defaults["top_p"],
        "top_k": args.top_k,
        "max_tokens": args.max_tokens,
        "mode": "thinking" if args.thinking else "non-thinking",
    }
    if args.samples > 1 and settings["temperature"] == 0:
        raise SystemExit("--samples above 1 at temperature 0 produces the same answer every time")

    seeds = read_jsonl(Path(args.seeds))
    if args.limit:
        seeds = seeds[:args.limit]
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    done = already_done(out_path) if args.resume else set()
    if done:
        print(f"resuming: {len(done)} answer(s) already in {out_path}")
    elif out_path.exists() and not args.resume:
        raise SystemExit(f"{out_path} exists; pass --resume to append to it, or choose another name")

    jobs = [(row, s) for row in seeds for s in range(args.samples)
            if f"{row['id']}#{s}" not in done]
    if not jobs:
        print("nothing to do: every prompt already has its answers")
        return

    endpoint = args.base_url.rstrip("/") + "/chat/completions"
    print(f"teacher: {args.model} at {endpoint}")
    print(f"settings: {json.dumps(settings)}")
    print(f"{len(jobs)} request(s), {args.concurrency} in flight")

    write_lock = threading.Lock()
    counters = {"prompt_tokens": 0, "completion_tokens": 0, "ok": 0, "failed": 0}
    power = PowerSampler()
    power.start()
    started = time.time()

    def one(job: tuple[dict, int]) -> None:
        row, sample = job
        messages = []
        if args.system:
            messages.append({"role": "system", "content": args.system})
        messages.append({"role": "user", "content": row["prompt"]})
        payload = {
            "model": args.model,
            "messages": messages,
            "temperature": settings["temperature"],
            "top_p": settings["top_p"],
            "max_tokens": settings["max_tokens"],
        }
        began = time.time()
        try:
            body = post_json(endpoint, payload, args.api_key, args.timeout)
            answer = (body["choices"][0]["message"]["content"] or "").strip()
            usage = body.get("usage", {}) or {}
        except (RuntimeError, KeyError, IndexError) as exc:
            with write_lock:
                counters["failed"] += 1
                print(f"  FAILED {row['id']}#{sample}: {exc}", file=sys.stderr)
            return

        record = {
            "id": row["id"],
            "sample": sample,
            "category": row.get("category"),
            "prompt": row["prompt"],
            "completion": answer,
            "teacher": {"id": args.teacher_id, "served_as": args.model, "quant": args.quant},
            "settings": settings,
            "usage": {
                "prompt_tokens": usage.get("prompt_tokens"),
                "completion_tokens": usage.get("completion_tokens"),
            },
            "seconds": round(time.time() - began, 2),
        }
        if "answer" in row:
            record["answer"] = row["answer"]  # carried through for the verifier

        with write_lock:
            with out_path.open("a", encoding="utf-8") as handle:
                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
                handle.flush()
            counters["ok"] += 1
            counters["prompt_tokens"] += usage.get("prompt_tokens") or 0
            counters["completion_tokens"] += usage.get("completion_tokens") or 0
            if counters["ok"] % 25 == 0:
                rate = counters["ok"] / max(time.time() - started, 1e-6)
                print(f"  {counters['ok']}/{len(jobs)} answers, {rate:.1f} per second")

    with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        futures = [pool.submit(one, job) for job in jobs]
        for future in as_completed(futures):
            future.result()

    elapsed = time.time() - started
    power.stop()
    mean_watts = power.mean()

    print()
    print(f"{counters['ok']} answered, {counters['failed']} failed, {elapsed / 60:.1f} minutes")
    print(f"prompt tokens {counters['prompt_tokens']}, completion tokens {counters['completion_tokens']}")
    print(f"mean accelerator power: {mean_watts if mean_watts is not None else 'not reported here'}")
    print(f"written to {out_path}")

    if args.labbook:
        cost = distillog.build_cost(
            prompt_tokens=counters["prompt_tokens"],
            completion_tokens=counters["completion_tokens"],
            seconds=elapsed,
            mean_watts=mean_watts,
        )
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/generate-teacher-data",
            stage="generate",
            teacher={"id": args.teacher_id, "served_as": args.model, "quant": args.quant,
                     "base_url": args.base_url},
            student=None,
            dataset={
                "seeds": args.seeds,
                "seeds_sha256": distillog.file_sha256(args.seeds),
                "out": str(out_path),
                "out_sha256": distillog.file_sha256(out_path),
                "prompts": len(seeds),
                "answers": counters["ok"],
                "failed": counters["failed"],
            },
            hyperparameters={**settings, "samples": args.samples, "concurrency": args.concurrency},
            seed=None,
            cost=cost,
            config_path=__file__,
            notes=None,
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()