#!/usr/bin/env python3
"""Convert a merged fish-native model dir (model.pth + dual_ar config) back to the HF fish_qwen3_omni
layout (safetensors shards + original index/config/tokenizer/codec) for sglang-omni / vLLM-omni.
Exact inverse of _remap_fish_qwen3_omni_keys, replayed over the ORIGINAL s2-pro index.

  python scripts/export_hf.py --src checkpoints/s2-pro --merged checkpoints/s2pro-egy-v2 --dst checkpoints/s2pro-egy-v2-hf
"""
import argparse
import json
import shutil
from collections import OrderedDict, defaultdict
from pathlib import Path

import torch
from safetensors.torch import save_file

ap = argparse.ArgumentParser()
ap.add_argument("--src", default="checkpoints/s2-pro")
ap.add_argument("--merged", required=True)
ap.add_argument("--dst", required=True)
a = ap.parse_args()
SRC, MERGED, DST = Path(a.src), Path(a.merged), Path(a.dst)
DST.mkdir(parents=True, exist_ok=True)


def forward_map(hf_key: str) -> str:
    if hf_key.startswith("text_model.model."):
        return hf_key[len("text_model.model."):]
    if hf_key.startswith("audio_decoder."):
        suffix = hf_key[len("audio_decoder."):]
        return suffix if suffix.startswith("codebook_embeddings.") else "fast_" + suffix
    return hf_key


index = json.loads((SRC / "model.safetensors.index.json").read_text())
weight_map = index["weight_map"]
inv = {}
for hf_key in weight_map:
    ik = forward_map(hf_key)
    assert ik not in inv, f"collision: {ik}"
    inv[ik] = hf_key

sd = torch.load(MERGED / "model.pth", map_location="cpu", mmap=True, weights_only=True)
if "state_dict" in sd:
    sd = sd["state_dict"]
sd = {k[len("model."):] if k.startswith("model.") else k: v for k, v in sd.items() if "lora" not in k}
missing_in_index = [k for k in sd if k not in inv]
missing_in_model = [k for k in inv if k not in sd]
print(f"index keys {len(weight_map)} | merged keys {len(sd)} | not in index {len(missing_in_index)} | missing from merged {len(missing_in_model)}")
assert not missing_in_index, missing_in_index[:5]
assert not missing_in_model, missing_in_model[:5]

shards = defaultdict(OrderedDict)
for ik, t in sd.items():
    shards[weight_map[inv[ik]]][inv[ik]] = t.to(torch.bfloat16).contiguous()
total = 0
for name, tensors in sorted(shards.items()):
    total += sum(t.numel() * t.element_size() for t in tensors.values())
    save_file(tensors, str(DST / name), metadata={"format": "pt"})
    print(f"wrote {name} ({len(tensors)} tensors)")
(DST / "model.safetensors.index.json").write_text(json.dumps({"metadata": {"total_size": total}, "weight_map": weight_map}, indent=2))
for f in ["config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja", "LICENSE.md", "codec.pth"]:
    if (SRC / f).exists():
        shutil.copy2(SRC / f, DST / f)
print(f"EXPORT_DONE {DST} total_size={total/1e9:.2f}GB")