"""Generate a format-following supervised fine-tuning dataset, and a held-out task file to score it.

Purpose: build a few hundred examples that all answer in one fixed output format, split them
         into a training set, a validation set and a held-out task file in the Part 10
         harness's shape, and write them in both the TRL and the mlx-lm layouts so either
         training path can read them. Every answer is derived from the structured fields
         that produced the question, so the labels are correct by construction rather than
         by an author's memory.
Platform: all (standard library only; the optional teacher step needs a served model)
Minimum memory: 8 GB on the machine running the optional teacher; this script needs very little
Assumes: Python 3.10 or newer. With --teacher-url, an OpenAI-compatible endpoint is reachable:
         llama-server from Part 6, the gateway from Part 9, or anything else speaking that API.
         The teacher is used only to rewrite the question text into more natural prose; the
         answers are never generated, because a label you did not check is not a label.

Usage: python3 make-format-dataset.py --out-dir . --count 320 --seed 0
       python3 make-format-dataset.py --print-template > my-format.json
       python3 make-format-dataset.py --template my-format.json --out-dir .
       python3 make-format-dataset.py --out-dir . \
           --teacher-url http://127.0.0.1:8080/v1 --teacher-model local/chat --teacher-count 80

Writes, under --out-dir:
  data/train.jsonl, data/valid.jsonl          TRL conversational prompt-completion
  data-mlx/{train,valid,test}.jsonl           mlx-lm completions
  format-tasks.json                           held-out tasks in the Part 10 task-file shape
  format-dataset.json                         everything generated, with its provenance
"""
from __future__ import annotations

import argparse
import json
import random
import sys
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# The built-in template. Everything about the format lives here, so replacing the
# format means replacing this object rather than editing the generator: run with
# --print-template, edit the copy, and pass it back with --template.
# ---------------------------------------------------------------------------
TEMPLATE: dict[str, Any] = {
    "name": "ticket-triage",
    "system": (
        "You triage incoming reports for a small operations team. "
        "Answer only in the required format, with no preamble and no closing remark."
    ),
    "instruction": (
        "Triage this report. Reply with exactly three lines and nothing else:\n"
        "Summary: one line, under twelve words\n"
        "Severity: one of low, medium or high\n"
        "Action: one imperative sentence"
    ),
    "answer_format": "Summary: {summary}\nSeverity: {severity}\nAction: {action}",
    "must_contain": ["Summary:", "Severity:", "Action:"],
    "must_not_contain": ["```", "Sure,", "Here is", "Certainly"],
    "max_words": 40,
    "reporters": [
        "A customer", "The overnight operator", "A colleague in support",
        "The monitoring system", "A developer on the platform team", "The duty manager",
    ],
    "systems": [
        {"name": "the billing API", "short": "billing"},
        {"name": "the document index", "short": "the index"},
        {"name": "the model gateway", "short": "the gateway"},
        {"name": "the nightly backup job", "short": "the backup job"},
        {"name": "the internal wiki", "short": "the wiki"},
        {"name": "the report generator", "short": "the report generator"},
    ],
    "symptoms": [
        {
            "text": "is returning errors for roughly one request in ten",
            "summary": "{short} failing intermittently",
            "severity": "high",
            "action": "Page the on-call engineer and check the error rate for {short}.",
        },
        {
            "text": "is unreachable from every machine we have tried",
            "summary": "{short} completely unreachable",
            "severity": "high",
            "action": "Page the on-call engineer and confirm whether {short} is running.",
        },
        {
            "text": "answers correctly but takes several times longer than usual",
            "summary": "{short} responding slowly",
            "severity": "medium",
            "action": "Check load and recent changes on {short} before escalating.",
        },
        {
            "text": "logged one failure that did not repeat on retry",
            "summary": "single transient failure in {short}",
            "severity": "low",
            "action": "Record the failure and watch {short} for a repeat.",
        },
        {
            "text": "has a spelling mistake on one of its pages",
            "summary": "cosmetic text error in {short}",
            "severity": "low",
            "action": "Open a low-priority ticket to correct the text in {short}.",
        },
        {
            "text": "returned a result that looks wrong but has not been confirmed",
            "summary": "possible incorrect output from {short}",
            "severity": "medium",
            "action": "Reproduce the case against {short} and confirm before escalating.",
        },
        {
            "text": "is filling the disk faster than expected",
            "summary": "{short} consuming disk quickly",
            "severity": "medium",
            "action": "Check free space and the retention settings for {short}.",
        },
        {
            "text": "stopped without an error message and had to be restarted by hand",
            "summary": "{short} exited silently",
            "severity": "high",
            "action": "Page the on-call engineer and collect the logs from {short}.",
        },
    ],
    "contexts": [
        "since this morning's deploy",
        "for the last two hours",
        "intermittently since the weekend",
        "starting a few minutes ago",
        "every night this week",
    ],
    # Deliberately hard cases. A dataset without them trains a model that answers
    # confidently when the input does not support an answer.
    "edge_cases": [
        {
            "input": "Someone said something is broken. No other detail was given.",
            "summary": "report lacks any detail to triage",
            "severity": "low",
            "action": "Ask the reporter which system failed and when.",
        },
        {
            "input": "The dashboard is green and a colleague says everything looks fine today.",
            "summary": "no fault reported",
            "severity": "low",
            "action": "Close the report with no action.",
        },
        {
            "input": "Two reports arrived at once: the wiki is slow, and the backup job exited silently.",
            "summary": "two faults reported together",
            "severity": "high",
            "action": "Split the report and handle the backup job first.",
        },
        {
            "input": "A customer asks when the next release is. Nothing appears to be broken.",
            "summary": "question rather than a fault report",
            "severity": "low",
            "action": "Forward the question to the release owner.",
        },
    ],
}


def post_json(url: str, payload: dict, api_key: str | None, timeout: int) -> dict:
    """One OpenAI-compatible request. Same shape as Part 10's harness uses."""
    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")[:300]
        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 rewrite_with_teacher(text: str, args, timeout: int = 120) -> str:
    """Ask a served model to make one report read like a person wrote it.

    The teacher never sees the answer and never writes one. It rewrites the question so
    that the training inputs are not all the same sentence with the nouns swapped, which
    is the one job here where a wrong output costs nothing: a badly rewritten question is
    still a question, whereas a badly written label is a wrong label.
    """
    payload = {
        "model": args.teacher_model,
        "messages": [
            {"role": "system", "content":
                "Rewrite the user's message as a short, natural report from a colleague. "
                "Keep every fact, keep it under 40 words, change nothing about what "
                "happened, and reply with the rewritten report only."},
            {"role": "user", "content": text},
        ],
        "temperature": 0.8,
        "max_tokens": 120,
    }
    data = post_json(f"{args.teacher_url.rstrip('/')}/chat/completions",
                     payload, args.teacher_key, timeout)
    content = (data.get("choices") or [{}])[0].get("message", {}).get("content", "")
    content = " ".join(str(content).split())
    return content or text


def build_examples(template: dict[str, Any], count: int, rng: random.Random) -> list[dict[str, str]]:
    """Every combination is drawn without replacement, so no two examples are identical."""
    combos = [
        (reporter, system, symptom, context)
        for reporter in template["reporters"]
        for system in template["systems"]
        for symptom in template["symptoms"]
        for context in template["contexts"]
    ]
    rng.shuffle(combos)
    edges = list(template["edge_cases"])
    wanted_edges = min(len(edges), max(1, count // 12))
    body = count - wanted_edges
    if body > len(combos):
        raise SystemExit(
            f"--count {count} needs {body} distinct combinations but the template only has "
            f"{len(combos)}. Add reporters, systems, symptoms or contexts, or ask for fewer."
        )

    examples: list[dict[str, str]] = []
    for reporter, system, symptom, context in combos[:body]:
        short = system["short"]
        examples.append({
            "input": f"{reporter} reports: {system['name']} {symptom['text']} {context}.",
            "summary": symptom["summary"].format(short=short),
            "severity": symptom["severity"],
            "action": symptom["action"].format(short=short),
            "kind": "generated",
        })
    for edge in edges[:wanted_edges]:
        examples.append({
            "input": edge["input"], "summary": edge["summary"],
            "severity": edge["severity"], "action": edge["action"], "kind": "edge",
        })
    rng.shuffle(examples)
    return examples


def answer_of(template: dict[str, Any], example: dict[str, str]) -> str:
    return template["answer_format"].format(
        summary=example["summary"], severity=example["severity"], action=example["action"])


def prompt_of(template: dict[str, Any], example: dict[str, str]) -> str:
    return f"{example['input']}\n\n{template['instruction']}"


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    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")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--out-dir", default=".", help="directory to write the dataset into")
    parser.add_argument("--count", type=int, default=320, help="total examples before splitting")
    parser.add_argument("--valid-fraction", type=float, default=0.1)
    parser.add_argument("--task-count", type=int, default=24,
                        help="held-out examples turned into a Part 10 task file")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--template", default=None, help="a JSON template from --print-template")
    parser.add_argument("--print-template", action="store_true",
                        help="print the built-in template and exit, so you can edit a copy")
    parser.add_argument("--teacher-url", default=None,
                        help="OpenAI-compatible base URL, e.g. http://127.0.0.1:8080/v1")
    parser.add_argument("--teacher-model", default=None, help="model name the endpoint answers to")
    parser.add_argument("--teacher-key", default=None, help="API key, if the endpoint needs one")
    parser.add_argument("--teacher-count", type=int, default=0,
                        help="how many inputs the teacher should rewrite; 0 disables it")
    args = parser.parse_args()

    if args.print_template:
        json.dump(TEMPLATE, sys.stdout, indent=2, ensure_ascii=False)
        print()
        return

    template = TEMPLATE
    if args.template:
        template = json.loads(Path(args.template).read_text(encoding="utf-8"))
        for key in ("system", "instruction", "answer_format", "reporters", "systems",
                    "symptoms", "contexts", "edge_cases"):
            if key not in template:
                raise SystemExit(f"{args.template} is missing the required key {key!r}")

    if args.teacher_count and not (args.teacher_url and args.teacher_model):
        raise SystemExit("--teacher-count needs both --teacher-url and --teacher-model")

    rng = random.Random(args.seed)
    examples = build_examples(template, args.count, rng)

    rewritten = 0
    for example in examples[: args.teacher_count]:
        if example["kind"] == "edge":
            continue  # the edge cases are the point; leave their wording alone
        try:
            example["input"] = rewrite_with_teacher(example["input"], args)
            example["kind"] = "generated-rewritten"
            rewritten += 1
        except RuntimeError as exc:
            print(f"teacher call failed, keeping the original wording: {exc}")
            break

    n_tasks = min(args.task_count, len(examples) // 4)
    held_out = examples[:n_tasks]
    remaining = examples[n_tasks:]
    n_valid = max(1, int(len(remaining) * args.valid_fraction))
    valid, train = remaining[:n_valid], remaining[n_valid:]

    out = Path(args.out_dir)
    system_message = {"role": "system", "content": template["system"]}

    def trl_rows(rows: list[dict[str, str]]) -> list[dict[str, Any]]:
        return [{
            "prompt": [system_message, {"role": "user", "content": prompt_of(template, r)}],
            "completion": [{"role": "assistant", "content": answer_of(template, r)}],
        } for r in rows]

    def mlx_rows(rows: list[dict[str, str]]) -> list[dict[str, Any]]:
        return [{"prompt": prompt_of(template, r), "completion": answer_of(template, r)}
                for r in rows]

    write_jsonl(out / "data" / "train.jsonl", trl_rows(train))
    write_jsonl(out / "data" / "valid.jsonl", trl_rows(valid))
    write_jsonl(out / "data-mlx" / "train.jsonl", mlx_rows(train))
    write_jsonl(out / "data-mlx" / "valid.jsonl", mlx_rows(valid))
    write_jsonl(out / "data-mlx" / "test.jsonl", mlx_rows(held_out))

    tasks = {
        "name": f"{template['name']}-format",
        "version": 1,
        "note": ("Held-out format tasks generated alongside the training data and never trained "
                 "on. Scored by Part 10's run-eval.py: the deterministic checks alone say whether "
                 "the output contract is being honoured, and the rubric lets a judge grade the "
                 "content of the three lines."),
        "settings": {"temperature": 0.0, "top_p": 1.0, "seed": 7, "max_tokens": 256,
                     "comment": "Fixed. The base model and the fine-tune are scored at these "
                                "same settings or the comparison means nothing."},
        "categories": {"format": "Does it produce the three required lines and nothing else?"},
        "tasks": [
            {
                "id": f"f{index + 1:02d}",
                "category": "format",
                "prompt": prompt_of(template, example),
                "reference": answer_of(template, example),
                "rubric": ("Exactly three lines, in the order Summary, Severity, Action, with no "
                           "preamble and no closing sentence. Severity is one of low, medium or "
                           "high and matches the seriousness of the report. The summary names "
                           "what failed. The action is a single imperative sentence."),
                "must_contain": template["must_contain"],
                "must_not_contain": template["must_not_contain"],
                "max_words": template["max_words"],
            }
            for index, example in enumerate(held_out)
        ],
    }
    (out / "format-tasks.json").write_text(
        json.dumps(tasks, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    (out / "format-dataset.json").write_text(json.dumps({
        "template": template["name"],
        "seed": args.seed,
        "count": len(examples),
        "teacher": {"url": args.teacher_url, "model": args.teacher_model,
                    "rewritten": rewritten} if rewritten else None,
        "splits": {"train": len(train), "validation": len(valid), "held_out_tasks": len(held_out)},
        "examples": examples,
    }, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    print(f"template:        {template['name']}")
    print(f"train:           {len(train)} example(s) -> {out / 'data' / 'train.jsonl'}")
    print(f"validation:      {len(valid)} example(s) -> {out / 'data' / 'valid.jsonl'}")
    print(f"held-out tasks:  {len(held_out)} -> {out / 'format-tasks.json'}")
    print(f"mlx-lm layout:   {out / 'data-mlx'}")
    if rewritten:
        print(f"teacher rewrote {rewritten} input(s); the answers were derived, not generated")
    print("\nRead ten training examples before you train on three hundred.")
    print("Then run decontaminate.py against your own Part 10 task file.")


if __name__ == "__main__":
    main()