#!/usr/bin/env python3
"""Build the two Swift1.5 campaign corpora (calibration + KLD).

calib.jsonl — 10M tokens of agent-shaped RAW text (45% code, 30% prose,
25% tool-call JSON), one {"prompt": text} per line, pushed to
/v1/completions with max_tokens=1 during the calibration run (the Hessian
tap is of the prompt tokens; docs/MOE-W4.md step 2). Code and prose mix
the repo checkout with HF dataset parquets (Frank's call 2026-10-06:
local material is only ~1.6M unique tokens; MOE-W4 wants 10M):
nampdn-ai/tiny-codes (code) and Salesforce/wikitext wikitext-103-raw
(prose). Tool calls are synthesized from commands harvested from the
repo's own scripts plus a fleet-ops set, with argument-shaped variants.

kld.jsonl — 1.8M tokens, one {"prompt", "score_from", "source"} per
line; each prompt is a multi-turn conversation rendered through the
Froggeric v22.5 template with a >= --kld-min-prompt (8192) token target,
and the last 256 tokens (the tail of the assistant answer) are the
scored span: score_from = n - 256. At 8k-token prompts the 1.8M corpus
is ~220 prompts = ~56k scored positions ~= 25 GB of record, matching
the published Flash-Next KLD scale (83k positions, MOE-W4.md). The KLD
reference run and every candidate run consume the SAME file (TOOLS.md
--kld-record / --kld-ref). The KLD corpus is LOCAL material only, so the
file-level disjointness rule (fnv1a64(path) % 5 == 0 -> KLD pool) is
untouched by the HF additions.

Deps: tokenizers, jinja2, pyarrow (venv /home/dank/venvs/radcorpus).
Usage: make-calib-corpus.py --repo ROOT --tokenizer-json PATH
    --template PATH --out-dir DIR [--hf-code-parquet F,F]
    [--hf-prose-parquet F,F] [--calib-tokens 10000000]
    [--kld-tokens 1800000]
"""
import argparse
import json
import os
import random
import re
import sys

from jinja2 import Environment, FileSystemLoader

_HERE = os.path.dirname(os.path.abspath(__file__)) if "__file__" in globals() else os.getcwd()
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

import qwen38_common as qc  # noqa: E402 - sibling module, same directory

CODE_EXT = (".sh", ".py", ".jinja", ".recipe", ".env", ".example",
            ".yaml", ".yml", ".bats", ".service", ".container", ".toml")
JSONL_EXTS = (".jsonl", ".json")
DOC_DIR = "serverAppConfigs-docs"
FRAGMENT_TOKENS = 2400   # what one fragment may contribute
CLASS_MIX = {"code": .45, "prose": .30, "tools": .25}
HF_CAP = {"code": 40000, "prose": 60000, "tools": 0}
POOL_CYCLES = 40         # how many times a class pool may repeat before giving up
KLD_GUARD = 5000         # rendered conversations tried before declaring the pool short

FLEET_CMDS = [
    "systemctl status frank-archon --no-pager",
    "journalctl -u frank-archon -n 100 --no-pager",
    "podman ps --format '{{.Names}} {{.Status}}'",
    "podman logs --tail 50 frank-archon",
    "curl -s http://192.168.0.235:7770/health",
    "curl -s http://192.168.0.235:7770/metrics | grep radiance_",
    "df -h /home/dank | tail -1",
    "free -g",
    "sudo systemctl restart frank-archon",
    "cd /home/dank/serverAppConfigs && git pull --ff-only",
    "./scripts/deploy.sh dankserver --service frank-archon",
    "./tests/run.sh 2>&1 | tail -5",
    "jq . /home/dank/models/rad/index.json | head -40",
    "python3 -c \"import json; print(json.load(open('config.json'))['num_hidden_layers'])\"",
    "ip neigh show 192.168.0.235",
    "amdgpu_top --device-info 2>/dev/null | head -20",
    "grep -E 'MemAvailable' /proc/meminfo",
    "sha256sum model-00001-of-00131.safetensors",
    "git log --oneline -10",
    "awk '{s+=$1} END {print s}' shard_sizes.txt",
]
TOOL_NAMES = ["bash", "read", "grep", "web_search", "edit", "write"]
WORKDIRS = ["/home/dank", "/home/dank/serverAppConfigs", "/home/dank/models",
            "/etc/systemd/system", "/var/log"]
DESCS = ["check service state", "tail the engine ledger", "list containers",
          "query health", "read metrics", "check disk", "check memory",
          "restart the endpoint", "sync the checkout", "run deployer",
          "run the test seams", "inspect the index", "probe the peer"]


def fnv(s: str) -> int:
    h = 0xcbf29ce484222325
    for b in s.encode():
        h = ((h ^ b) * 0x100000001b3) & 0xFFFFFFFFFFFFFFFF
    return h


def harvest_cmds(repo: str, rng: random.Random) -> list[str]:
    cmds = list(FLEET_CMDS)
    pat = re.compile(r"^\s*(?:sudo\s+)?([a-z][\w./-]*(?:\s+\S+){1,6})\s*$")
    for root, dirs, files in os.walk(repo):
        dirs[:] = [d for d in dirs if not d.startswith(".")]
        for f in files:
            if not f.endswith(".sh"):
                continue
            try:
                for line in open(os.path.join(root, f), encoding="utf-8",
                                 errors="replace"):
                    m = pat.match(line.rstrip())
                    if m and len(m.group(1)) > 12:
                        cmds.append("ssh dank@192.168.0.235 '"
                                    + m.group(1)[:90] + "'")
            except OSError:
                pass
    rng.shuffle(cmds)
    return cmds[:4000]


def tool_call(rng: random.Random, cmds: list[str]) -> str:
    """One agent turn: 2-4 tool calls with argument-shaped JSON."""
    parts = []
    for _ in range(rng.randint(2, 4)):
        args: dict = {"command": rng.choice(cmds)}
        if rng.random() < 0.4:
            args["cwd"] = rng.choice(WORKDIRS)
        if rng.random() < 0.3:
            args["timeout"] = rng.choice([10, 30, 60, 300, 600])
        if rng.random() < 0.5:
            args["description"] = rng.choice(DESCS)
        parts.append('<tool_call>{"name": "%s", "arguments": %s}<|im_end|>'
                     % (rng.choice(TOOL_NAMES),
                        json.dumps(args, ensure_ascii=False)))
    return "\n".join(parts)


def load_parquet_texts(path: str, kind: str, rng: random.Random,
                       cap: int) -> list[str]:
    import pyarrow.parquet as pq
    d = pq.read_table(path).to_pydict()
    n = len(next(iter(d.values())))
    idx = range(n) if n <= cap else rng.sample(range(n), cap)
    out = []
    for i in idx:
        if kind == "code":
            if "content" in d:
                frag = d["content"][i]
            else:
                frag = ("// %s\n### Instruction:\n%s\n### Response:\n%s"
                        % (d.get("file_name", [""])[i], d["prompt"][i],
                           d["code"][i] or d["response"][i]))
        else:
            frag = d["text"][i]
        if frag and frag.strip():
            out.append(frag[:8000])
    rng.shuffle(out)
    return out[:cap]


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--repo", required=True)
    ap.add_argument("--tokenizer-json", required=True)
    ap.add_argument("--template", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--calib-tokens", type=int, default=10_000_000)
    ap.add_argument("--kld-tokens", type=int, default=1_800_000)
    ap.add_argument("--chunk-tokens", type=int, default=8192)
    ap.add_argument("--score-tail", type=int, default=256)
    ap.add_argument("--kld-min-prompt", type=int, default=8192)
    ap.add_argument("--hf-code-parquet", default="")
    ap.add_argument("--hf-prose-parquet", default="")
    ap.add_argument("--seed", type=int, default=20261006)
    return ap.parse_args()


def load_template(template: str) -> tuple[Environment, str]:
    """The chat template as the engine will render it, plus its short name."""
    env = Environment(loader=FileSystemLoader(os.path.dirname(template)),
                      trim_blocks=True, lstrip_blocks=True)
    env.filters["tojson"] = lambda v: json.dumps(v, ensure_ascii=False)
    env.filters["raise_exception"] = lambda m: (_ for _ in ()).throw(
        ValueError(m))
    return env, os.path.basename(template)


def classify(repo: str) -> dict[str, list[str]]:
    """Walk the checkout into the three source classes the mix draws from."""
    code, prose, jsonl = [], [], []
    for root, dirs, files in os.walk(repo):
        dirs[:] = [d for d in dirs if not d.startswith(".")]
        for f in sorted(files):
            p = os.path.relpath(os.path.join(root, f), repo)
            if DOC_DIR in p and f.endswith(".md"):
                prose.append(p)
            elif p.startswith(DOC_DIR):
                pass
            elif f.endswith(CODE_EXT):
                code.append(p)
            elif f.endswith(JSONL_EXTS) and f != "package-lock.json":
                jsonl.append(p)
    return {"code": code, "prose": prose, "tools": jsonl}


def split_pools(classes: dict[str, list[str]]) -> dict[str, tuple[list, list]]:
    """File-level split: bucket 0 -> KLD pool, everything else -> calibration."""
    return {cls: ([p for p in lst if fnv(p) % 5 == 0],
                  [p for p in lst if fnv(p) % 5 != 0])
            for cls, lst in classes.items()}


def jsonl_content(raw: str) -> str:
    """The content fields of a content-keyed JSONL, or "" for anything else.

    A .json in this repo is not necessarily a content-keyed JSONL: the substring
    test passes on any line that merely mentions "content", and a list-shaped
    document raises TypeError — both used to crash the whole build. Skip what
    does not fit instead.
    """
    lines = []
    for line in raw.splitlines():
        if not line.strip() or "content" not in line:
            continue
        try:
            rec = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(rec, dict) and isinstance(rec.get("content"), str):
            lines.append(rec["content"])
    return "\n".join(lines)


def fragment_text(repo: str, rel: str, tok) -> str:
    """One file as a fragment: its first FRAGMENT_TOKENS tokens, decoded."""
    raw = open(os.path.join(repo, rel), encoding="utf-8",
               errors="replace").read()
    if rel.endswith(JSONL_EXTS):
        raw = jsonl_content(raw)
    ids = tok.encode(raw, add_special_tokens=False).ids
    return tok.decode(ids[:FRAGMENT_TOKENS], skip_special_tokens=True)


def build_frags(a: argparse.Namespace, tok,
                rng: random.Random) -> dict[str, dict[str, list[str]]]:
    """Both fragment pools per class, seeded so the same seed rebuilds the corpus."""
    pools = split_pools(classify(a.repo))
    cmds = harvest_cmds(a.repo, rng)
    frags: dict[str, dict[str, list[str]]] = {}
    for cls, (kld_files, cal_files) in pools.items():
        if cls == "tools":
            frags[cls] = {"kld": [tool_call(rng, cmds) for _ in range(600)],
                          "cal": [tool_call(rng, cmds) for _ in range(40000)]}
        else:
            frags[cls] = {"kld": [fragment_text(a.repo, p, tok) for p in kld_files],
                          "cal": [fragment_text(a.repo, p, tok) for p in cal_files]}
    # per-class fragment caps: prose rows are short (~66 tok) so need many;
    # codeparrot files are long (~2.6k tok) so a small cap suffices
    for cls, spec in (("code", a.hf_code_parquet),
                      ("prose", a.hf_prose_parquet)):
        for path in [p for p in spec.split(",") if p]:
            frags[cls]["cal"].extend(
                load_parquet_texts(path, cls, rng, cap=HF_CAP[cls]))
    for cls in frags:
        rng.shuffle(frags[cls]["cal"])
    return frags


def pack_class(cls: str, pool: list[str], want: int, tok,
               chunk_tokens: int) -> list[dict]:
    """Greedily pack `pool` into chunk_tokens-sized calibration prompts."""
    if not pool:
        return []
    rows: list[dict] = []
    # Count every fragment once: the loop cycles the pool up to POOL_CYCLES
    # times, and re-tokenizing a 2.4k-token fragment on each pass dominated the
    # build's wall time.
    counts = [qc.n_tokens(tok, f) for f in pool]
    got = 0
    i = 0
    buf: list[str] = []
    nbuf = 0
    while got < want:
        j = i % len(pool)
        frag = pool[j]
        i += 1
        n = counts[j]
        if nbuf + n > chunk_tokens and buf:
            rows.append({"prompt": "\n".join(buf), "cls": cls})
            got += nbuf
            buf, nbuf = [], 0
            if got >= want:
                break
        buf.append(frag)
        nbuf += n
        if i > len(pool) * POOL_CYCLES:
            break  # pool exhausted with repeats; stop
    if buf and got < want:
        rows.append({"prompt": "\n".join(buf), "cls": cls})
        got += nbuf
    print(f"calib {cls}: {got} tokens, {i} fragments")
    return rows


def emit_calib(a: argparse.Namespace, frags, tok,
               rng: random.Random) -> tuple[int, list[dict]]:
    """calib.jsonl rows: the class mix, shuffled, prompts only."""
    rows: list[dict] = []
    for cls, frac in CLASS_MIX.items():
        rows.extend(pack_class(cls, frags[cls]["cal"], int(a.calib_tokens * frac),
                               tok, a.chunk_tokens))
    rng.shuffle(rows)
    return len(rows), [{"prompt": r["prompt"]} for r in rows]


def emit_kld(a: argparse.Namespace, frags, tok, rng: random.Random,
             env: Environment, tname: str) -> tuple[int, list[dict], int]:
    """kld.jsonl rows: multi-turn conversations, each with a scored tail."""
    template = env.get_template(tname)
    rows: list[dict] = []
    total = scored_total = 0
    classes = ["code", "prose", "tools"]
    cursor = {c: 0 for c in classes}
    guard = 0
    while total < a.kld_tokens and guard < KLD_GUARD:
        guard += 1
        msgs, srcs = [], []
        for t in range(rng.choice([2, 3, 4])):
            cls = classes[(len(rows) + t) % 3]
            pool = frags[cls]["kld"]
            if not pool:
                continue
            frag = pool[cursor[cls] % len(pool)]
            cursor[cls] += 1
            msgs.append({"role": "user",
                         "content": "Summarize and explain this "
                                    f"engineering material:\n{frag}"})
            msgs.append({"role": "assistant",
                         "content": frag[:4000]})
            srcs.append(cls)
        rendered = template.render(messages=msgs, add_generation_prompt=False)
        n = qc.n_tokens(tok, rendered)
        if n < a.kld_min_prompt:
            continue
        rows.append({"prompt": rendered,
                     "score_from": n - a.score_tail,
                     "source": "+".join(sorted(set(srcs)))})
        total += n
        scored_total += a.score_tail
    return len(rows), rows, scored_total


def manifest(a: argparse.Namespace, tname: str, nc: int, nk: int,
             scored: int) -> dict:
    return {
        "seed": a.seed, "template": tname,
        "tokenizer": os.path.basename(a.tokenizer_json),
        "calib_prompts": nc, "kld_prompts": nk,
        "kld_scored_positions": scored,
        "kld_score_tail": a.score_tail,
        "disjoint_file_pools": "local KLD material only; HF code feeds calib",
        "calib_mix": dict(CLASS_MIX),
        "hf_code": a.hf_code_parquet, "hf_prose": a.hf_prose_parquet,
    }


def main() -> None:
    a = parse_args()
    rng = random.Random(a.seed)
    tok = qc.load_tokenizer(a.tokenizer_json)
    env, tname = load_template(a.template)

    frags = build_frags(a, tok, rng)
    os.makedirs(a.out_dir, exist_ok=True)
    nc, calib_rows = emit_calib(a, frags, tok, rng)
    qc.write_jsonl(os.path.join(a.out_dir, "calib.jsonl"), calib_rows)
    nk, kld_rows, scored = emit_kld(a, frags, tok, rng, env, tname)
    qc.write_jsonl(os.path.join(a.out_dir, "kld.jsonl"), kld_rows)
    man = manifest(a, tname, nc, nk, scored)
    qc.write_json(os.path.join(a.out_dir, "manifest.json"), man)
    print(json.dumps(man, indent=1))


if __name__ == "__main__":
    main()