#!/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

from jinja2 import Environment, FileSystemLoader
from tokenizers import Tokenizer

CODE_EXT = (".sh", ".py", ".jinja", ".recipe", ".env", ".example",
            ".yaml", ".yml", ".bats", ".service", ".container", ".toml")
JSONL_EXTS = (".jsonl", ".json")
DOC_DIR = "serverAppConfigs-docs"

FLEET_CMDS = [
    "systemctl status frank-primary --no-pager",
    "journalctl -u frank-primary -n 100 --no-pager",
    "podman ps --format '{{.Names}} {{.Status}}'",
    "podman logs --tail 50 frank-primary",
    "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-primary",
    "cd /home/dank/serverAppConfigs && git pull --ff-only",
    "./scripts/deploy.sh dankserver --service frank-primary",
    "./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 main() -> None:
    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)
    a = ap.parse_args()

    rng = random.Random(a.seed)
    tokenizer = Tokenizer.from_file(a.tokenizer_json)
    env = Environment(loader=FileSystemLoader(os.path.dirname(a.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))
    tname = os.path.basename(a.template)

    code, prose, jsonl = [], [], []
    for root, dirs, files in os.walk(a.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), a.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)
    # file-level split: bucket 0 -> KLD pool, 1..4 -> calibration pool
    pools: dict[str, tuple[list, list]] = {}
    for cls, lst in (("code", code), ("prose", prose), ("tools", jsonl)):
        pools[cls] = ([p for p in lst if fnv(p) % 5 == 0],
                      [p for p in lst if fnv(p) % 5 != 0])

    def text_of(p: str) -> str:
        path = os.path.join(a.repo, p)
        raw = open(path, encoding="utf-8", errors="replace").read()
        if p.endswith(JSONL_EXTS):
            lines = [json.loads(l)["content"] for l in raw.splitlines()
                     if l.strip() and "content" in l]
            raw = "\n".join(lines)
        toks = tokenizer.encode(raw, add_special_tokens=False).ids
        return tokenizer.decode(toks[:2400], skip_special_tokens=True)

    cmds = harvest_cmds(a.repo, rng)
    # fragment pools per class, corpus-split first
    frags: dict[str, dict[str, list[str]]] = {}
    for cls in pools:
        kld_files, cal_files = pools[cls]
        fr = {}
        src = {"kld": kld_files, "cal": cal_files}
        if cls == "tools":
            fr["kld"] = [tool_call(rng, cmds) for _ in range(600)]
            fr["cal"] = [tool_call(rng, cmds) for _ in range(40000)]
        else:
            fr["kld"] = [text_of(p) for p in src["kld"]]
            fr["cal"] = [text_of(p) for p in src["cal"]]
        frags[cls] = fr
    # 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
    hfcap = {"code": 40000, "prose": 60000, "tools": 0}
    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=hfcap[cls]))
    for cls in frags:
        rng.shuffle(frags[cls]["cal"])

    def toks(s: str) -> int:
        return len(tokenizer.encode(s, add_special_tokens=False).ids)

    def emit_calib() -> tuple[int, list[dict]]:
        targets = {"code": .45, "prose": .30, "tools": .25}
        rows: list[dict] = []
        for cls, frac in targets.items():
            want = int(a.calib_tokens * frac)
            pool = frags[cls]["cal"]
            if not pool:
                continue
            got, i, buf, nbuf = 0, 0, [], 0
            while got < want:
                frag = pool[i % len(pool)]
                i += 1
                n = toks(frag)
                if nbuf + n > a.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) * 40:
                    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")
        rng.shuffle(rows)
        return len(rows), [{k: v for k, v in r.items() if k == "prompt"}
                           for r in rows]

    def emit_kld() -> tuple[int, list[dict]]:
        rows, total = [], 0
        classes = ["code", "prose", "tools"]
        i = {c: 0 for c in classes}
        guard = 0
        while total < a.kld_tokens and guard < 5000:
            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[i[cls] % len(pool)]
                i[cls] += 2
                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 = env.get_template(tname).render(
                messages=msgs, add_generation_prompt=False)
            n = toks(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
        return len(rows), rows

    os.makedirs(a.out_dir, exist_ok=True)
    nc, rows = emit_calib()
    with open(os.path.join(a.out_dir, "calib.jsonl"), "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    nk, krows = emit_kld()
    with open(os.path.join(a.out_dir, "kld.jsonl"), "w") as f:
        for r in krows:
            f.write(json.dumps(r) + "\n")
    scored = sum(len(tokenizer.encode(r["prompt"], add_special_tokens=False).ids)
                 - r["score_from"] for r in krows)
    manifest = {
        "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": {"code": .45, "prose": .30, "tools": .25},
        "hf_code": a.hf_code_parquet, "hf_prose": a.hf_prose_parquet,
    }
    with open(os.path.join(a.out_dir, "manifest.json"), "w") as f:
        json.dump(manifest, f, indent=1)
    print(json.dumps(manifest, indent=1))


if __name__ == "__main__":
    main()