#!/usr/bin/env python3
"""Shared plumbing for the qwen38 campaign corpus tools.

The three tools move the same two shapes — a JSONL of prompt records and a
one-artifact JSON manifest — and they all tokenize with the SAME tokenizer. That
last part is not a detail: a KLD reference is only meaningful against the exact
token sequence the corpus recorded, so `n_tokens` exists once, here, and every
tool counts the same way. A window that decode/encode does not round-trip is a
refusal, never a warning.

Loaded as a sibling module. Each tool puts its own directory on `sys.path`
first, so `python3 scripts/<tool>.py` and the driver's container mount
(`runpy.run_path('/work/<tool>.py')`, `-v scripts:/work:ro`) both resolve it —
run the tools by path, not piped through stdin.
"""
import contextlib
import json
import os
import sys
import tempfile


def refuse(msg: str) -> None:
    """Exit 2 with a REFUSED line — the code argparse uses for unusable input."""
    print(f"REFUSED: {msg}", file=sys.stderr)
    sys.exit(2)


def load_tokenizer(path: str):
    try:
        from tokenizers import Tokenizer
    except ImportError as exc:  # the corpus tools carry their own venv
        refuse(f"tokenizers is not importable by {sys.executable} — run the "
               f"corpus venv's python (on dankserver: "
               f"/home/dank/venvs/radcorpus/bin/python) or set "
               f"FRANK_QWEN38_TOKENIZER: {exc}")
    if not os.path.isfile(path):
        refuse(f"no tokenizer at {path}")
    return Tokenizer.from_file(path)


def n_tokens(tok, text: str) -> int:
    """The corpus's own count: no special tokens added, prompt side only."""
    return len(tok.encode(text, add_special_tokens=False).ids)


def iter_jsonl(path: str):
    """Yield `(lineno, record)` for every non-blank line; refuse on a bad line.

    Field validation stays with the caller: the three corpora carry different
    fields (`prompt` alone, `prompt`+`cls`, `prompt`+`score_from`+`source`), and
    a record missing the one the tool needs is that tool's refusal to make.
    """
    with open(path, encoding="utf-8", errors="replace") as fh:
        for lineno, line in enumerate(fh, 1):
            if not line.strip():
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError as exc:
                refuse(f"{path}:{lineno} is not JSON: {exc}")
            yield lineno, rec


def write_jsonl(path: str, rows) -> None:
    with open(path, "w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")


@contextlib.contextmanager
def atomic_writer(path: str):
    """Write `path` through a temporary sibling and rename on success.

    An interrupted run then leaves no half corpus behind, and the temporary name
    keeps the one-artifact-one-name law the driver applies to a score report.
    """
    d = os.path.dirname(os.path.abspath(path)) or "."
    fd, tmp = tempfile.mkstemp(dir=d, prefix="." + os.path.basename(path) + "-")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as fh:
            yield fh
        os.replace(tmp, path)
    except BaseException:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise
    os.chmod(path, 0o644)


def write_json(path: str, obj) -> None:
    with atomic_writer(path) as fh:
        json.dump(obj, fh, indent=1)
        fh.write("\n")