"""Build a domain fine-tuning dataset from your own material, with a review gate before training.

Purpose: turn work you already did, or drafts a local teacher wrote from your documents, into the
         train, validation and held-out splits a fine-tune needs. Two modes: --from-pairs takes
         input and output pairs you wrote or reviewed, and --from-docs asks a served model to
         draft candidates from your documents into a review file that is not training data until
         you have marked each one accepted.
Platform: all (standard library only; --from-docs needs a served model over the OpenAI-compatible
          API, which may be on any track or another machine)
Minimum memory: 8 GB on the machine serving the teacher; this script needs very little
Assumes: Python 3.10 or newer. For --from-docs, an OpenAI-compatible endpoint at --teacher-url
         and a directory of .txt or .md files. Everything written here carries whatever licence
         and confidentiality your source material carries, so redact before you generate rather
         than before you publish.

Usage: python3 make-domain-dataset.py --from-pairs my-pairs.jsonl --out-dir . \
           --system "You answer questions about our deployment runbook." \
           --instruction "Answer in at most four sentences and cite the section."

       python3 make-domain-dataset.py --from-docs ~/runbooks --out-dir . \
           --teacher-url http://127.0.0.1:8080/v1 --teacher-model qwen3-8b \
           --per-chunk 2 --max-chunks 60
       # then read drafts.jsonl, set "accepted" on every record, and:
       python3 make-domain-dataset.py --from-pairs drafts.jsonl --out-dir .

Writes, under --out-dir:
  data/train.jsonl, data/valid.jsonl      TRL conversational prompt-completion
  data-mlx/{train,valid,test}.jsonl       mlx-lm completions
  domain-tasks.json                       held-out tasks in the Part 10 task-file shape
  drafts.jsonl                            --from-docs only: candidates awaiting your review
"""
from __future__ import annotations

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

DEFAULT_SYSTEM = "You answer questions about this team's own documents, briefly and exactly."
DEFAULT_INSTRUCTION = (
    "Answer in at most four sentences. If the documents do not contain the answer, "
    "say so in one sentence instead of guessing."
)
DRAFT_SYSTEM = (
    "You write training examples. Given a passage, produce question and answer pairs that a "
    "colleague could plausibly ask and that the passage fully answers. Never use information "
    "that is not in the passage. Reply with one JSON object per line and nothing else, each "
    'of the form {"input": "...", "output": "..."}.'
)


def post_json(url: str, payload: dict, api_key: str | None, 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")[: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 chunk_text(text: str, words_per_chunk: int) -> list[str]:
    """Split on blank lines, then pack paragraphs up to a word budget.

    Paragraph boundaries rather than a fixed window, because a question drafted from half a
    sentence is a question whose answer is not in the passage.
    """
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    chunks, current, count = [], [], 0
    for paragraph in paragraphs:
        words = len(paragraph.split())
        if current and count + words > words_per_chunk:
            chunks.append("\n\n".join(current))
            current, count = [], 0
        current.append(paragraph)
        count += words
    if current:
        chunks.append("\n\n".join(current))
    return chunks


def draft_from_docs(args) -> list[dict[str, Any]]:
    """Ask the teacher for candidates. Nothing here is training data yet."""
    docs_dir = Path(args.from_docs).expanduser()
    paths = sorted(p for p in docs_dir.rglob("*") if p.suffix.lower() in {".txt", ".md"})
    if not paths:
        raise SystemExit(f"no .txt or .md files under {docs_dir}")
    print(f"{len(paths)} document(s) under {docs_dir}")

    drafts: list[dict[str, Any]] = []
    chunk_index = 0
    for path in paths:
        for chunk in chunk_text(path.read_text(encoding="utf-8", errors="replace"), args.words_per_chunk):
            if args.max_chunks and chunk_index >= args.max_chunks:
                break
            chunk_index += 1
            payload = {
                "model": args.teacher_model,
                "messages": [
                    {"role": "system", "content": DRAFT_SYSTEM},
                    {"role": "user", "content":
                        f"Write {args.per_chunk} question and answer pair(s) from this passage. "
                        f"Answers must follow this instruction: {args.instruction}\n\n"
                        f"Passage:\n{chunk}"},
                ],
                "temperature": 0.7,
                "max_tokens": 600,
            }
            try:
                body = post_json(f"{args.teacher_url.rstrip('/')}/chat/completions",
                                 payload, args.teacher_key, args.timeout)
            except RuntimeError as exc:
                print(f"teacher call failed on chunk {chunk_index}: {exc}")
                break
            content = (body.get("choices") or [{}])[0].get("message", {}).get("content", "") or ""
            for line in content.splitlines():
                line = line.strip().strip("`")
                if not line.startswith("{"):
                    continue
                try:
                    item = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if not (isinstance(item.get("input"), str) and isinstance(item.get("output"), str)):
                    continue
                drafts.append({
                    "id": f"d{len(drafts) + 1:04d}",
                    "source": str(path.relative_to(docs_dir)),
                    "input": item["input"].strip(),
                    "output": item["output"].strip(),
                    "accepted": None,
                })
            print(f"  chunk {chunk_index} from {path.name}: {len(drafts)} draft(s) so far")
    return drafts


def load_pairs(path: Path, accept_unreviewed: bool) -> list[dict[str, Any]]:
    """Read input/output pairs, honouring the review gate."""
    pairs, unreviewed, rejected = [], 0, 0
    with path.open("r", encoding="utf-8") as handle:
        for number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
            except json.JSONDecodeError as exc:
                raise SystemExit(f"{path}:{number}: not valid JSON ({exc.msg})") from exc
            if not (isinstance(item.get("input"), str) and isinstance(item.get("output"), str)):
                raise SystemExit(f"{path}:{number}: every record needs string 'input' and 'output'")
            accepted = item.get("accepted")
            if accepted is False:
                rejected += 1
                continue
            if accepted is None:
                unreviewed += 1
                if not accept_unreviewed:
                    continue
            pairs.append(item)
    if rejected:
        print(f"{rejected} record(s) marked accepted: false, skipped")
    if unreviewed:
        state = "included" if accept_unreviewed else "skipped"
        print(f"{unreviewed} record(s) have no 'accepted' field, {state}")
        if not accept_unreviewed:
            print("  Set \"accepted\": true on the ones you have read, or pass "
                  "--accept-unreviewed if you wrote every pair yourself.")
    return pairs


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])
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--from-pairs", help="JSON Lines with input, output and accepted")
    source.add_argument("--from-docs", help="directory of .txt or .md files for the teacher to read")
    parser.add_argument("--out-dir", default=".")
    parser.add_argument("--system", default=DEFAULT_SYSTEM,
                        help="system prompt; use the same one at serving time or the fine-tune "
                             "is being asked to generalise across a change you never trained")
    parser.add_argument("--instruction", default=DEFAULT_INSTRUCTION,
                        help="appended to every question, and the contract the answers honour")
    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("--max-words", type=int, default=120,
                        help="deterministic length check written into the task file")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--accept-unreviewed", action="store_true",
                        help="use records with no 'accepted' field; only for pairs you wrote")
    parser.add_argument("--teacher-url", default=None)
    parser.add_argument("--teacher-model", default=None)
    parser.add_argument("--teacher-key", default=None)
    parser.add_argument("--per-chunk", type=int, default=2)
    parser.add_argument("--words-per-chunk", type=int, default=250)
    parser.add_argument("--max-chunks", type=int, default=0, help="0 means every chunk")
    parser.add_argument("--timeout", type=int, default=300)
    args = parser.parse_args()

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    if args.from_docs:
        if not (args.teacher_url and args.teacher_model):
            raise SystemExit("--from-docs needs --teacher-url and --teacher-model")
        drafts = draft_from_docs(args)
        if not drafts:
            raise SystemExit("the teacher produced no usable drafts; check the endpoint and model")
        write_jsonl(out / "drafts.jsonl", drafts)
        print(f"\n{len(drafts)} draft(s) written to {out / 'drafts.jsonl'}")
        print("These are not training data yet. Read every one, set \"accepted\" to true or")
        print("false, and then re-run with --from-pairs drafts.jsonl. The rejections are the")
        print("most informative part: a pattern in them is usually a defect in the prompt to")
        print("the teacher rather than in the teacher.")
        return

    pairs = load_pairs(Path(args.from_pairs), args.accept_unreviewed)
    if len(pairs) < 40:
        print(f"WARNING: only {len(pairs)} usable pair(s). A domain behaviour change usually "
              "wants low thousands; a format change wants a few hundred. Expect a small effect.")
    if not pairs:
        raise SystemExit("no usable pairs; nothing to write")

    rng = random.Random(args.seed)
    rng.shuffle(pairs)

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

    system_message = {"role": "system", "content": args.system}
    prompt_of = lambda item: f"{item['input']}\n\n{args.instruction}"  # noqa: E731

    write_jsonl(out / "data" / "train.jsonl", [{
        "prompt": [system_message, {"role": "user", "content": prompt_of(p)}],
        "completion": [{"role": "assistant", "content": p["output"]}],
    } for p in train])
    write_jsonl(out / "data" / "valid.jsonl", [{
        "prompt": [system_message, {"role": "user", "content": prompt_of(p)}],
        "completion": [{"role": "assistant", "content": p["output"]}],
    } for p in valid])
    for name, rows in (("train", train), ("valid", valid), ("test", held_out)):
        write_jsonl(out / "data-mlx" / f"{name}.jsonl",
                    [{"prompt": prompt_of(p), "completion": p["output"]} for p in rows])

    tasks = {
        "name": "domain-holdout",
        "version": 1,
        "note": ("Held-out examples from the same source as the training data, never trained on. "
                 "Score the base model and the fine-tune on this file at the same settings, and "
                 "score both on your own Part 10 task file as well: this one shows the gain and "
                 "that one shows the cost."),
        "settings": {"temperature": 0.0, "top_p": 1.0, "seed": 7, "max_tokens": 512,
                     "comment": "Fixed for both models. Changing them makes a new baseline."},
        "categories": {"domain": "Does it answer this domain's questions in the required form?"},
        "tasks": [{
            "id": f"d{index + 1:02d}",
            "category": "domain",
            "prompt": prompt_of(item),
            "reference": item["output"],
            "rubric": (f"Answers the question from the domain material, honouring this "
                       f"instruction: {args.instruction} A confident answer to something the "
                       f"material does not cover scores 1."),
            "max_words": args.max_words,
        } for index, item in enumerate(held_out)],
    }
    (out / "domain-tasks.json").write_text(
        json.dumps(tasks, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    print(f"train:           {len(train)} -> {out / 'data' / 'train.jsonl'}")
    print(f"validation:      {len(valid)} -> {out / 'data' / 'valid.jsonl'}")
    print(f"held-out tasks:  {len(held_out)} -> {out / 'domain-tasks.json'}")
    print(f"mlx-lm layout:   {out / 'data-mlx'}")
    print("\nNext: run decontaminate.py against both this task file and your own Part 10 file.")


if __name__ == "__main__":
    main()