#!/usr/bin/env python3
"""Build the two Swift1.5 campaign corpora from the repo checkout.

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).

kld.jsonl — 1.8M tokens, one {"prompt", "score_from", "source"} per
line; each prompt is a two-turn conversation rendered through the
Froggeric v22.5 template, and the last 256 tokens (the tail of the
assistant answer) are the scored span: score_from = n - 256. The KLD
reference run and every candidate run consume the SAME file (TOOLS.md
--kld-record / --kld-ref).

Disjointness is at file granularity: fnv1a64(relative path) % 5 == 0
goes to the KLD pool, the rest to the calibration pool; synthesized tool
calls partition by template id. The manifest records the token counts
per source class and the empty intersection.

Content sources: shell/python/jinja/recipe/config files (code),
serverAppConfigs-docs markdown (prose), and synthesized agent tool
calls — JSON envelopes whose commands are harvested from the repo's own
scripts (ssh-prefixed) plus a vendored fleet-ops command set.

Deps: tokenizers, jinja2 (see /home/dank/venvs/radcorpus on dankserver).
The tokenizer loads from the checkpoint's tokenizer.json; the template
is the deployed Froggeric v22.5 file. Usage:

  make-calib-corpus.py --repo ROOT --tokenizer-json PATH --template PATH
      --out-dir DIR [--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",
    "cat /proc/meminfo | grep -E 'MemAvailable|MemAvailable'",
    "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"]
TOOL_JSON = ('<tool_call>{"name": "%s", "arguments": {"command": "%s"}}'
             '<|im_end|>\n')


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:
    return TOOL_JSON % (rng.choice(TOOL_NAMES),
                        rng.choice(cmds).replace('"', "'"))


def render(env: Environment, tokenizer: Tokenizer, template_name: str,
           messages: list) -> tuple[str, int]:
    out = env.get_template(template_name).render(
        messages=messages, add_generation_prompt=False)
    return out, len(tokenizer.encode(out, add_special_tokens=False).ids)


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("--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)

    # 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":
            cmds = harvest_cmds(a.repo, rng)
            src["kld"] = [f"synth:{i}" for i in range(0, len(cmds), 2)]
            src["cal"] = [f"synth:{i}" for i in range(1, len(cmds), 2)]
        fr["kld"] = [tool_call(rng, src["kld"][i::1]) if cls == "tools"
                     else text_of(p) for i, p in enumerate(src["kld"])]
        fr["cal"] = [tool_call(rng, src["cal"][i::1]) if cls == "tools"
                     else text_of(p) for i, p in enumerate(src["cal"])]
        frags[cls] = fr

    def emit_calib() -> tuple[int, list[str]]:
        targets = {"code": .45, "prose": .30, "tools": .25}
        rows, used = [], []
        for cls, frac in targets.items():
            want = int(a.calib_tokens * frac)
            pool = frags[cls]["cal"]
            rng.shuffle(pool)
            got, i, buf = 0, 0, []
            nbuf = 0
            while got < want and i < len(pool) * 3:
                frag = pool[i % len(pool)]
                n = len(tokenizer.encode(frag, add_special_tokens=False).ids)
                if nbuf + n > a.chunk_tokens and buf:
                    rows.append({"prompt": "\n".join(buf)})
                    used.extend(set())
                    got += nbuf
                    buf, nbuf = [], 0
                buf.append(frag)
                nbuf += n
                i += 1
            if buf:
                rows.append({"prompt": "\n".join(buf)})
                got += nbuf
        rng.shuffle(rows)
        return len(rows), rows

    def emit_kld() -> tuple[int, list[str]]:
        rows, total = [], 0
        classes = ["code", "prose", "tools"]
        i = {c: 0 for c in classes}
        while total < a.kld_tokens:
            msgs, srcs, ntok = [], [], 0
            turns = rng.choice([2, 3])
            for t in range(turns):
                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, n = render(env, tokenizer, tname, msgs)
            if n <= a.score_tail:
                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"][max(0, 0):], 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": True,
        "calib_mix": {"code": .45, "prose": .30, "tools": .25},
    }
    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()