#!/usr/bin/env python3
"""K1: write offload/launch/configs/D-kv6-1m-k1mem.env, D-kv6-pack.env and D-kv6-pack-1m.env from D-cal1-kv6.env (remote Engram
+ remote DSpark capacity, KV 6 GiB/rank) plus the K1 overlay, and the same three on the async production candidate
D-kv6-async-process-nap-lazy.env (S8R2 FIX-2: D-kv6-async-1m-k1mem.env, D-kv6-async-pack.env, D-kv6-async-pack-1m.env; the
GB10 plan runs these, the sync ones stay as the control). Every K1 item is its own variable with a default that a
start can override (empty = off), e.g.  K1_TF_ALIAS= offload/launch/start_exp.sh offload/launch/configs/D-kv6-async-pack-1m.env
  python3 offload/kvpack/make_configs.py          # write
  python3 offload/kvpack/make_configs.py --check  # the committed configs equal a fresh generation
"""
import re
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
CFG = ROOT / "offload/launch/configs"
SYNC_BASE = "D-cal1-kv6.env"
ASYNC_BASE = "D-kv6-async-process-nap-lazy.env"   # window 4's production candidate; R3's TP2 + async package uses the same settings
NODROP_SPEC = ('{"method":"dspark","num_speculative_tokens":5,"num_speculative_tokens_per_batch_size":[[1,2,5],[3,4,4]],'
               '"disable_eagle_block_drop":true}')


def base_parts(base: str) -> tuple[str, str, str]:
    text = (CFG / base).read_text()
    get = lambda k: re.search(rf'^{k}="(.*)"$', text, re.M).group(1)  # noqa: E731
    env = get("EXP_DOCKER_ENV")
    if "VLLM_INDEXER_PREFILL_BUFFER_FACTOR" in env:
        raise SystemExit("base config already sets the indexer factor")
    return get("OVERLAYS"), get("EXTRA_DOCKER_ARGS"), env


def k1_overlays() -> str:
    return subprocess.run([str(ROOT / "offload/kvpack/stage_overlay.sh"), "overlays"], check=True,
                          capture_output=True, text=True).stdout.strip()


def render(label: str, header: str, max_len: int | None, kv_pair: str, base: str) -> str:
    overlays, extra, env = base_parts(base)
    lines = [
        f"# {header}",
        f"# Generated by offload/kvpack/make_configs.py from {base} + the K1 overlay (offload/kvpack, MANIFEST.json).",
        "# Before the first start on a node: offload/kvpack/stage_overlay.sh check spark sparkworker (image base == overlay base).",
        "# K1 switches (override at start; empty = off). Without the overlay mounted none of them exists in the image.",
        f"K1_KV_PAIR=${{K1_KV_PAIR-{kv_pair}}}            # KV page packing: 3 = 0060 (three 64-state pages), 1 = 0057; empty = v5 layout",
        "K1_DROP_MTP=${K1_DROP_MTP-1}            # 0076b: no [4096, 20480] bf16 MTP residual buffer (160 MiB/rank, never read by DSpark)",
        "K1_TF_TRANSIENT=${K1_TF_TRANSIENT-1}    # 0076d: TensorFold prompt scratch per MoE call instead of 1,014 MiB/rank held",
        "K1_TF_ALIAS=${K1_TF_ALIAS-1}            # 0076e (with 0076d): down partials in the rotated rows' storage, call peak 1,014 -> 750 MiB",
        "K1_INDEXER_FACTOR=${K1_INDEXER_FACTOR-4} # indexer prefill K-gather workspace factor (image default 8); empty = 8",
        "K1_MEMPROBE=${K1_MEMPROBE-}             # 0076a: per-step memory JSONL prefix /root/.cache/dsv41-memprobe-<run>; empty = off",
        "# memprobe opens <prefix>_r<rank>.jsonl and creates no directory: the prefix must sit directly in /root/.cache",
        "# (= ~/.cache/dsv41-serve on each node, mounted by the launcher), or the worker dies at its first dummy step.",
        'case "$K1_MEMPROBE" in "") ;; /root/.cache/*/*|/root/.cache/) echo "K1_MEMPROBE=$K1_MEMPROBE: use /root/.cache/<name> (no subdirectory)" >&2; return 1 ;; /root/.cache/*) ;; *) echo "K1_MEMPROBE=$K1_MEMPROBE: must start with /root/.cache/" >&2; return 1 ;; esac',
        "K1_NODROP=${K1_NODROP-}                # disable_eagle_block_drop: NO effect in our serve (S8 N-4: with swa_bounded_replay the",
        "#                                         flagged SWA group is not prefix-cacheable, so the drop never applies); a switch only",
        f"LABEL={label}",
        "KV_BYTES_OVERRIDE=${K1_KV_BYTES-6442450944}   # 6 GiB/rank; K1_KV_BYTES=1610612736 (1.5 GiB) = pool-pressure run R1b",
    ]
    if max_len:
        lines.append(f"MAX_LEN_OVERRIDE={max_len}")
    lines += [
        f"if [ -n \"$K1_NODROP\" ]; then SPEC_OVERRIDE='{NODROP_SPEC}'; fi",
        f'OVERLAYS="{overlays} {k1_overlays()}"',
        f'EXTRA_DOCKER_ARGS="{extra}"',
        f'EXP_DOCKER_ENV="{env}'
        '${K1_KV_PAIR:+ VLLM_DSV41_KV_PAIR=$K1_KV_PAIR}'
        '${K1_DROP_MTP:+ VLLM_DSV41_DROP_MTP_BUFFER=$K1_DROP_MTP}'
        '${K1_TF_TRANSIENT:+ VLLM_EXL3_TF_PROMPT_TRANSIENT=$K1_TF_TRANSIENT}'
        '${K1_TF_ALIAS:+ VLLM_EXL3_TF_PROMPT_ALIAS=$K1_TF_ALIAS}'
        '${K1_INDEXER_FACTOR:+ VLLM_INDEXER_PREFILL_BUFFER_FACTOR=$K1_INDEXER_FACTOR}'
        '${K1_MEMPROBE:+ VLLM_DSV41_MEMPROBE=$K1_MEMPROBE VLLM_DSV41_MEMPROBE_LAYERS=2,14,20}"',
    ]
    return "\n".join(lines) + "\n"


CONFIGS = {
    "D-kv6-1m-k1mem.env": ("D-kv6-1m-k1mem", "D-kv6 with max-model-len 1,048,576 (the model's max_position_embeddings; YaRN x16 from 64K), "
                     "no KV packing, 0076b/d/e + indexer factor 4 (upstream: what makes 1M fit on GB10). Needs the x86host "
                     "worker started with --max-model-len 1048576. (D-kv6-1m.env is the coordinator's plain 1M run; not generated here.)", 1_048_576, "",
                     SYNC_BASE),
    "D-kv6-pack.env": ("D-kv6-pack", "D-kv6 (327,680) with KV page packing 0060 (VLLM_DSV41_KV_PAIR=3) + 0076b/d/e + factor 4; "
                       "predicted pool 2,502,600 tokens (v5: 1,604,926).", None, "3", SYNC_BASE),
    "D-kv6-pack-1m.env": ("D-kv6-pack-1m", "D-kv6-1m-k1mem with KV page packing 0060 + 0076b/d/e + factor 4; predicted pool "
                          "3,154,934 tokens = 3.01 requests of 1,048,576 (v5: 1,850,243). Needs the 1M x86host worker.",
                          1_048_576, "3", SYNC_BASE),
}
# S8R2 FIX-2: the same three rebased onto the async production candidate. Sourced as start_exp.sh sources them, they
# differ from the sync ones only in the base's five async variables (VLLM_DSPARK_REMOTE_ASYNC, _ASYNC_PROXY, _BACKOFF_MAX_MS,
# _ASYNC_SPIN_US, _ASYNC_NAP_US); headers and K1 switches are the same.
CONFIGS |= {name.replace("D-kv6-", "D-kv6-async-", 1): (label.replace("D-kv6-", "D-kv6-async-", 1), header, max_len, kv_pair,
                                                        ASYNC_BASE)
            for name, (label, header, max_len, kv_pair, _) in list(CONFIGS.items())}


def main() -> None:
    check = "--check" in sys.argv[1:]
    bad = []
    for name, (label, header, max_len, kv_pair, base) in CONFIGS.items():
        text = render(label, header, max_len, kv_pair, base)
        path = CFG / name
        if check:
            if not path.exists() or path.read_text() != text:
                bad.append(name)
        else:
            path.write_text(text)
            print("wrote", path.relative_to(ROOT))
    if bad:
        raise SystemExit("configs differ from a fresh generation: " + ", ".join(bad))
    if check:
        print(f"verified {len(CONFIGS)} configs")


if __name__ == "__main__":
    main()