#!/usr/bin/env python3
"""Export a serving bundle (DESIGN §9): merged text-only backbone (bf16 safetensors, HF layout of
Qwen3_5ForCausalLM without lm_head) + head.safetensors + judge_config.json + calibration.json + tokenizer.

python3 scripts/export.py --checkpoint checkpoints/s2_seed42/best --calibration checkpoints/s2_seed42/best/calibration.json \
    --out exports/jev-judge-qwen35-9b --name jev-judge-qwen35-9b --version 0.7.0
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys

import torch
from safetensors.torch import save_file

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

from jev_judge import __version__  # noqa: E402
from jev_judge.calibration import load as load_calibration  # noqa: E402
from jev_judge.checkpointing import load_checkpoint  # noqa: E402


def save_sharded(state: dict[str, torch.Tensor], out_dir: str, max_shard_bytes: int = 4 * 1024**3) -> None:
    shards: list[dict[str, torch.Tensor]] = [{}]
    sizes = [0]
    for k, v in state.items():
        nbytes = v.numel() * v.element_size()
        if sizes[-1] + nbytes > max_shard_bytes and shards[-1]:
            shards.append({}); sizes.append(0)
        shards[-1][k] = v
        sizes[-1] += nbytes
    n = len(shards)
    weight_map = {}
    for i, sh in enumerate(shards):
        name = f"model-{i+1:05d}-of-{n:05d}.safetensors"
        save_file({k: t.contiguous() for k, t in sh.items()}, os.path.join(out_dir, name), metadata={"format": "pt"})
        weight_map.update({k: name for k in sh})
    with open(os.path.join(out_dir, "model.safetensors.index.json"), "w") as f:
        json.dump({"metadata": {"total_size": sum(sizes)}, "weight_map": weight_map}, f, indent=2)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--calibration", required=True, help="calibration.json (fail-fast at serve time if missing)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--name", default="jev-judge")
    ap.add_argument("--version", default=__version__)
    ap.add_argument("--include-lm-head", action="store_true",
                    help="also ship the base model's lm_head.weight so the bundle can do ordinary AR generation")
    ap.add_argument("--mode", choices=["merged", "unmerged"], default="unmerged",
                    help="merged: LoRA folded into the shipped backbone (decision-only bundle); "
                         "unmerged: ship the *pristine* base backbone + adapter/ (dual-mode: AR generation is exactly the base model)")
    args = ap.parse_args()

    table = load_calibration(args.calibration)  # validates schema
    judge, cfg = load_checkpoint(args.checkpoint, device="cuda")
    lm = judge.lm
    os.makedirs(args.out, exist_ok=True)
    if args.mode == "unmerged":
        if not hasattr(lm, "peft_config"):
            raise SystemExit("checkpoint has no LoRA adapter; use --mode merged")
        lm.save_pretrained(os.path.join(args.out, "adapter"))  # peft adapter (fp32 master weights)
        lm = lm.unload()  # drop adapter layers -> pristine base weights
        args.include_lm_head = True  # dual-mode bundles always ship the vocab head
    elif hasattr(lm, "merge_and_unload"):
        lm = lm.merge_and_unload()
    lm.lm_head = None
    state = {k: v.detach().cpu() for k, v in lm.state_dict().items() if not k.startswith("lm_head")}  # dtypes as loaded (bf16 / fp32)
    if args.include_lm_head:
        # the judge dropped lm_head at load time; take the pristine rows from the base checkpoint
        import glob

        from safetensors import safe_open

        found = False
        for shard in sorted(glob.glob(os.path.join(cfg["base_model_path"], "*.safetensors"))):
            with safe_open(shard, framework="pt", device="cpu") as f:
                if "lm_head.weight" in f.keys():
                    state["lm_head.weight"] = f.get_tensor("lm_head.weight").to(torch.bfloat16)
                    found = True
                    break
        if not found:
            raise SystemExit("lm_head.weight not found in base checkpoint")
    save_sharded(state, args.out)
    # text-only config in Qwen3_5ForCausalLM layout
    tcfg = lm.config.to_dict()
    tcfg["architectures"] = ["Qwen3_5ForCausalLM"]
    with open(os.path.join(args.out, "config.json"), "w") as f:
        json.dump(tcfg, f, indent=2)
    save_file({k: v.detach().to(torch.float32).cpu().contiguous() for k, v in judge.head.state_dict().items()},
              os.path.join(args.out, "head.safetensors"))
    jc = judge.judge_config({"model_name": args.name, "model_version": args.version, "lm_head_included": bool(args.include_lm_head),
                             "weights_mode": args.mode, "adapter_subfolder": "adapter" if args.mode == "unmerged" else None,
                             "lora": cfg.get("lora"), "trained_stage": cfg.get("stage"),
                             "source_checkpoint": os.path.abspath(args.checkpoint), "softcap": judge.head.softcap})
    with open(os.path.join(args.out, "judge_config.json"), "w") as f:
        json.dump(jc, f, indent=2)
    shutil.copy(args.calibration, os.path.join(args.out, "calibration.json"))
    judge.tokenizer.save_pretrained(args.out)
    total = sum(v.numel() * v.element_size() for v in state.values())
    print(f"exported [{args.mode}] {len(state)} tensors ({total/1e9:.1f} GB) + head + calibration + tokenizer"
          f"{' + adapter/' if args.mode == 'unmerged' else ''} -> {args.out}")


if __name__ == "__main__":
    main()