#!/usr/bin/env python3
"""Build the Doc Quatrain training data: ten poems in, two JSONL files out.

    python3 build_corpus.py          # with the ten poems beside it
    python3 build_corpus.py <dir>    # or point it at them

    train.jsonl     8 examples
    holdout.jsonl   2 examples, never shown to the training run

Standard library only. Part eight of the series explains every choice made here.
"""
import json
import re
import sys
from pathlib import Path

DIR = Path(sys.argv[1] if len(sys.argv) > 1 else __file__).resolve()
DIR = DIR if DIR.is_dir() else DIR.parent

# The poems are output with no record of what was asked for, so the instruction side of
# each pair is written by hand. Note that none of these mentions rhyme or verse: the voice
# has to come from training rather than from the request.
PROMPTS = {
    "01": "Write a field note about backups that are never restored.",
    "02": "Write a field note about a stale DNS record and a resolver cache.",
    "03": "Write a field note about trusting a UPS battery you never load tested.",
    "04": "Write a field note about a TLS certificate that expired unnoticed.",
    "05": "Write a field note about dust, cooling, and a machine that quietly throttled.",
    "06": "Write a field note about debug logging that filled a disk.",
    "07": "Write a field note about the spare part you decided not to buy.",
    "08": "Write a field note about a firmware update that reset a configuration.",
    "09": "Write a field note about the two cables you never got around to labeling.",
    "10": "Write a field note about monitoring that failed silently.",
}

# Kept back so the finished model can be tested on subjects it has never seen. These two
# sit furthest from the rest, so passing the test needs actual generalization.
HOLDOUT = {"05", "10"}

# Kept short deliberately. A longer persona prompt would do the work the training is meant
# to do, leaving no way to tell which of the two produced the result.
SYSTEM = "You are Doc Quatrain. You write short technical field notes."

files = sorted(p for p in DIR.iterdir() if re.fullmatch(r"\d\d-.*\.md", p.name))
if len(files) != 10:
    raise SystemExit(f"expected 10 corpus files in {DIR}, found {len(files)}")

train, holdout = [], []

for path in files:
    key = path.name[:2]
    if key not in PROMPTS:
        raise SystemExit(f"no prompt written for {path.name}")

    # Drop the markdown heading: leaving it in teaches the model that a title is part of
    # the voice, and it then appears in everything the model generates.
    body = re.sub(r"\A#[^\n]*\n+", "", path.read_text(encoding="utf-8")).strip()

    example = {
        "messages": [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": PROMPTS[key]},
            {"role": "assistant", "content": body},
        ]
    }
    (holdout if key in HOLDOUT else train).append(example)


def write(name, rows):
    out = DIR / name
    # ensure_ascii=False keeps the text readable; compact separators match what the
    # JavaScript version of this script emitted, so the committed JSONL does not churn
    # just because the builder changed language. Neither choice matters to a trainer.
    out.write_text(
        "\n".join(json.dumps(r, ensure_ascii=False, separators=(",", ":")) for r in rows) + "\n",
        encoding="utf-8",
    )
    words = sum(len(r["messages"][2]["content"].split()) for r in rows)
    print(f"{name:<14} {len(rows):>2} examples  ~{words} words")


write("train.jsonl", train)
write("holdout.jsonl", holdout)