"""Write FLA kernel-choice config directories from recorded release choices.

FLA's gated-delta kernels pick block sizes by timing them in each process, so
two runtime processes can answer the Qwen3.5-based Decision 2.0 sizes
differently by rounding. vllm-project/semantic-router records each model's
released kernel choices in `src/model-runtime/vllm_sr_runtime/registry/
kernel_choices.json`. This script turns one device class of that file into
FLA's own config files (one directory per model, one `<kernel>.json` per
kernel), in the format FLA reads with `FLA_CONFIG_DIR=<dir>` and
`FLA_CACHE_MODE=full`, so every process runs the released configurations and
times none.

    python3 fla_kernel_choices.py kernel_choices.json /srv/decision-studio/fla-config \
      --model vllm-sr/Decision-2.0-Eos-0.8B=eos --model vllm-sr/Decision-2.0-Lux-9B=lux
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path


def fla_key_hash(key: list) -> str:
    """FLA's `AutotuneKey.key_hash` of a tuning key."""
    serialized = json.dumps(key, separators=(",", ":"), sort_keys=True)
    return hashlib.md5(serialized.encode(), usedforsecurity=False).hexdigest()


def write_model(choices: dict, directory: Path) -> int:
    directory.mkdir(parents=True, exist_ok=True)
    for stale in directory.glob("*.json"):
        stale.unlink()
    for name, entries in choices["kernels"].items():
        document = {
            "kernel_name": name,
            "autotune_entries": {
                fla_key_hash(entry["key"]): {"autotune_key": entry["key"], "config": entry["config"]}
                for entry in entries
            },
            "default_config": entries[0]["config"],
        }
        (directory / f"{name}.json").write_text(json.dumps(document, sort_keys=True), encoding="utf-8")
    return len(choices["kernels"])


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("kernel_choices")
    parser.add_argument("output")
    parser.add_argument("--device-class", default="rocm:gfx942")
    parser.add_argument("--fla", default="0.5.2", help="FLA version installed in the runtime image")
    parser.add_argument("--model", action="append", required=True, metavar="REPO=NAME")
    args = parser.parse_args(argv)
    recorded = json.loads(Path(args.kernel_choices).read_text(encoding="utf-8"))
    for item in args.model:
        repo, _, name = item.partition("=")
        entry = recorded.get(repo)
        choices = (entry or {}).get("choices", {}).get(args.device_class)
        if not choices or choices.get("fla") != args.fla:
            parser.error(f"{repo}: no {args.device_class} choices recorded with FLA {args.fla}")
        count = write_model(choices, Path(args.output) / name)
        print(f"{repo}@{entry.get('revision')} -> {name}: {count} kernels", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())