#!/usr/bin/env python3
"""Build the training data a draft head needs: the target model's own answers, and its features.

Purpose: a draft is trained to agree with one particular target model, so its training set
    is not human text but the target's own generations. This script has two stages. The
    "generate" stage sends your prompts through the served model and writes the answers as
    a ShareGPT-shaped JSONL file, which is the format the Medusa and EAGLE training
    repositories read. The "hidden-states" stage runs the same conversations back through
    the target with transformers and saves the second-to-top-layer hidden state per token,
    which is what an EAGLE-style draft regresses on. Run stage one on any track; run stage
    two only if the recipe you chose needs features.
Platform: generate: all (spark, strix, mac, nvidia; standard library only, talks to a
    server). hidden-states: spark, strix, nvidia and mac, and needs torch and transformers
    installed; it is the only stage that loads the weights locally.
Minimum memory: 8 GB for the generate stage. The hidden-states stage needs the target model
    resident: roughly 2 bytes per parameter at bfloat16, so about 17 GB for an 8B target and
    about 9 GB for a 4B one, plus the activations for one sequence.
Assumes: for generate, a served model reachable over plain HTTP at an OpenAI-compatible
    endpoint (Part 9's gateway, or llama-server, or vLLM). For hidden-states, torch and
    transformers in the active environment and the target's Hugging Face directory or id.
    Both stages append one record to the lab notebook through draftlog.py, which must sit
    next to this file.

Usage:
    # stage 1: the target's own answers to your prompts
    python3 make-draft-data.py generate --base-url http://127.0.0.1:8000/v1 \\
        --model my-fine-tune --prompts prompts.txt --out data/draft-data.jsonl \\
        --max-tokens 384 --temperature 0.0 --labbook labbook.md

    # stage 2: the features an EAGLE-style draft regresses on
    python3 make-draft-data.py hidden-states --target ./runs/my-fine-tune-merged \\
        --data data/draft-data.jsonl --out-dir data/hidden --limit 512 --labbook labbook.md

An API key, if the server needs one, is read from the environment variable named by
--api-key-env. No key is written to this file, to the output or to the notebook.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from urllib.parse import urlsplit

import draftlog

# A starter prompt set for readers who have not collected their own yet. Twelve prompts is
# far too few to train a usable draft; it is enough to prove the pipeline runs end to end,
# and the page says where the real prompts should come from.
STARTER_PROMPTS = [
    "Summarise what a key-value cache is and why it grows with the conversation.",
    "Rewrite this shell command with error handling: curl -s localhost:8080/v1/models",
    "Explain the difference between prefill and decode to a colleague.",
    "Write a Python function that returns the median of a list without imports.",
    "Give four fields that must accompany a tokens-per-second measurement.",
    "Describe continuous batching in plain language.",
    "What is an acceptance rate in speculative decoding?",
    "List three reasons a language model server refuses to start.",
    "Turn this into JSON with keys name and port: llama-server 8080, vllm 8000.",
    "Explain quantisation to somebody who knows what a floating point number is.",
    "Why can a smaller model that fits beat a larger one that does not?",
    "Write three sentences about what to record after a benchmark run.",
]


# ------------------------------------------------------------------------ HTTP


class ServerError(Exception):
    """The server did not answer in a way this script can use."""


def post_json(url: str, payload: dict, api_key: str, timeout: float) -> dict:
    """One non-streaming POST. Returns the decoded JSON body."""
    body = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(url, data=body, method="POST")
    request.add_header("Content-Type", "application/json")
    if api_key:
        request.add_header("Authorization", f"Bearer {api_key}")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:  # noqa: S310 - http only, checked below
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:200]
        raise ServerError(f"HTTP {exc.code} from {url}: {detail}") from exc
    except (urllib.error.URLError, OSError, json.JSONDecodeError) as exc:
        raise ServerError(f"{type(exc).__name__} talking to {url}: {exc}") from exc


# --------------------------------------------------------------------- stage 1


def read_prompts(path: str | None) -> list[str]:
    """One prompt per non-blank line, or the starter set."""
    if not path:
        return list(STARTER_PROMPTS)
    with open(path, encoding="utf-8") as handle:
        prompts = [line.strip() for line in handle if line.strip()]
    if not prompts:
        raise SystemExit(f"{path} contained no prompts")
    return prompts


def stage_generate(args: argparse.Namespace) -> int:
    """Collect the target model's own answers and write them as ShareGPT-shaped JSONL."""
    parts = urlsplit(args.base_url)
    if parts.scheme != "http":
        print("This tool speaks plain HTTP only; point it at a localhost or LAN endpoint.", file=sys.stderr)
        return 2
    if not parts.hostname:
        print(f"Could not read a host from --base-url {args.base_url!r}.", file=sys.stderr)
        return 2
    chat_url = f"{parts.scheme}://{parts.netloc}{parts.path.rstrip('/') or '/v1'}/chat/completions"
    api_key = os.environ.get(args.api_key_env, "")

    prompts = read_prompts(args.prompts)
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"==> generating {len(prompts)} answer(s) from {args.model} at {args.base_url}")
    if args.temperature > 0:
        print("    temperature is above zero: the draft will learn a sampled slice of the target's")
        print("    behaviour, which is usually what you want for coverage. Record the value.")

    kept = 0
    failed = 0
    started = time.perf_counter()
    with out_path.open("w", encoding="utf-8") as out:
        for i, prompt in enumerate(prompts):
            payload = {
                "model": args.model,
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": args.max_tokens,
                "temperature": args.temperature,
                "stream": False,
            }
            try:
                body = post_json(chat_url, payload, api_key, args.timeout)
            except ServerError as exc:
                failed += 1
                print(f"    [{i + 1}/{len(prompts)}] failed: {exc}", file=sys.stderr)
                continue
            choices = body.get("choices") or [{}]
            answer = ((choices[0].get("message") or {}).get("content") or "").strip()
            if not answer:
                failed += 1
                print(f"    [{i + 1}/{len(prompts)}] the server returned an empty answer", file=sys.stderr)
                continue
            out.write(json.dumps({
                "id": f"draft-{i:06d}",
                "conversations": [
                    {"from": "human", "value": prompt},
                    {"from": "gpt", "value": answer},
                ],
            }, ensure_ascii=False) + "\n")
            kept += 1
            if (i + 1) % 20 == 0:
                print(f"    [{i + 1}/{len(prompts)}] {kept} kept, {failed} failed")

    wall = time.perf_counter() - started
    print(f"==> wrote {kept} example(s) to {out_path} in {wall:.1f} s ({failed} failed)")

    draftlog.record(
        labbook=args.labbook,
        lab="part-17/make-draft-data/generate",
        model=args.model,
        dataset={
            "path": str(out_path),
            "sha256": draftlog.file_sha256(out_path),
            "examples": kept,
            "prompt_source": args.prompts or "built-in starter set",
        },
        hyperparameters={
            "stage": "generate",
            "max_tokens": args.max_tokens,
            "temperature": args.temperature,
            "base_url": args.base_url,
        },
        seed=args.seed,
        notes=f"{failed} prompt(s) failed or returned nothing",
    )
    print(f"    recorded the run in {args.labbook}")
    return 1 if kept == 0 else 0


# --------------------------------------------------------------------- stage 2


def stage_hidden_states(args: argparse.Namespace) -> int:
    """Save the target's second-to-top-layer hidden states for each conversation.

    An EAGLE-style draft does not predict tokens from tokens; it predicts the target's
    own features one step ahead, and it is trained against the features the target
    actually produced. That is why this stage exists and why it needs the weights
    locally: no serving API exposes hidden states.
    """
    try:
        import torch  # noqa: PLC0415 - optional heavy dependency, only this stage needs it
        from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415
    except ImportError as exc:
        print(f"This stage needs torch and transformers in the active environment: {exc}", file=sys.stderr)
        return 2

    data_path = Path(args.data)
    if not data_path.is_file():
        print(f"{data_path} does not exist; run the generate stage first.", file=sys.stderr)
        return 2
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else (
        "mps" if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available() else "cpu"
    )
    dtype = torch.bfloat16 if device == "cuda" else torch.float32
    print(f"==> loading {args.target} on {device} in {dtype}")

    tokenizer = AutoTokenizer.from_pretrained(args.target)
    model = AutoModelForCausalLM.from_pretrained(args.target, dtype=dtype)
    model.to(device)
    model.eval()

    index_path = out_dir / "index.jsonl"
    written = 0
    skipped = 0
    started = time.perf_counter()

    with data_path.open(encoding="utf-8") as handle, index_path.open("w", encoding="utf-8") as index:
        for line_no, line in enumerate(handle):
            if args.limit and written >= args.limit:
                break
            line = line.strip()
            if not line:
                continue
            example = json.loads(line)
            turns = example.get("conversations") or []
            messages = [
                {"role": "user" if t.get("from") == "human" else "assistant", "content": t.get("value", "")}
                for t in turns
            ]
            if len(messages) < 2:
                skipped += 1
                continue
            text = tokenizer.apply_chat_template(messages, tokenize=False)
            encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=args.max_length)
            input_ids = encoded["input_ids"].to(device)
            with torch.no_grad():
                out = model(input_ids, output_hidden_states=True)
            # hidden_states[0] is the embedding output and hidden_states[-1] the final layer,
            # so [-2] is the second-to-top layer the EAGLE paper describes.
            features = out.hidden_states[-2][0].to(torch.float16).cpu()
            target_path = out_dir / f"{example.get('id', f'row-{line_no:06d}')}.pt"
            torch.save({"input_ids": input_ids[0].cpu(), "hidden_states": features}, target_path)
            index.write(json.dumps({
                "id": example.get("id", f"row-{line_no:06d}"),
                "file": target_path.name,
                "tokens": int(input_ids.shape[1]),
                "hidden_size": int(features.shape[-1]),
            }) + "\n")
            written += 1
            if written % 25 == 0:
                print(f"    {written} example(s) written")

    wall = time.perf_counter() - started
    print(f"==> wrote {written} feature file(s) to {out_dir} in {wall:.1f} s ({skipped} skipped)")
    print("    These files are large. Check the size before you start a full run:")
    print(f"    du -sh {out_dir}")

    draftlog.record(
        labbook=args.labbook,
        lab="part-17/make-draft-data/hidden-states",
        model=args.target,
        dataset={
            "path": str(data_path),
            "sha256": draftlog.file_sha256(data_path),
            "examples": written,
            "features_dir": str(out_dir),
        },
        hyperparameters={
            "stage": "hidden-states",
            "layer": "second-to-top (hidden_states[-2])",
            "max_length": args.max_length,
            "device": device,
            "dtype": str(dtype),
        },
        seed=args.seed,
        notes=f"{skipped} row(s) skipped for having fewer than two turns",
    )
    print(f"    recorded the run in {args.labbook}")
    return 1 if written == 0 else 0


# ---------------------------------------------------------------------------- CLI


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--seed", type=int, default=0)
    sub = parser.add_subparsers(dest="stage", required=True)

    gen = sub.add_parser("generate", help="collect the target model's own answers through a served endpoint")
    gen.add_argument("--base-url", default="http://127.0.0.1:8000/v1")
    gen.add_argument("--model", required=True, help="model name the server reports at /v1/models")
    gen.add_argument("--prompts", default=None, help="file of prompts, one per line")
    gen.add_argument("--out", default="data/draft-data.jsonl")
    gen.add_argument("--max-tokens", type=int, default=384)
    gen.add_argument("--temperature", type=float, default=0.0)
    gen.add_argument("--timeout", type=float, default=300.0)
    gen.add_argument("--api-key-env", default="SPEC_API_KEY")

    hid = sub.add_parser("hidden-states", help="save the target's second-to-top-layer features per token")
    hid.add_argument("--target", required=True, help="Hugging Face directory or id of the target model")
    hid.add_argument("--data", default="data/draft-data.jsonl")
    hid.add_argument("--out-dir", default="data/hidden")
    hid.add_argument("--max-length", type=int, default=1024)
    hid.add_argument("--limit", type=int, default=0, help="stop after this many examples; 0 means all")

    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    if args.stage == "generate":
        return stage_generate(args)
    return stage_hidden_states(args)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))