#!/usr/bin/env python3
"""
export_hf.py: Export a chat-SFT checkpoint to a HuggingFace-style dir.

Output dir contains:
    model.safetensors       (weights, "model."-prefixed keys)
    config.json             (architecture, mask_vocab_size = 32010)
    generation_config.json  (sampling defaults for chat)
    tokenizer.json / tokenizer_config.json   (with ChatML + rainbow tokens)
    README.md               (minimal)

Usage:
    python3 export_hf.py \
        --checkpoint checkpoints_chat/step_3000.pt \
        --tokenizer data/no_robots_chatml/tokenizer \
        --output MetaDiffusion-150M-Chat/
"""

import argparse
import json
import shutil
import sys
from dataclasses import asdict
from pathlib import Path

import torch
from safetensors.torch import save_file

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from model import MetaDiffusionConfig  # noqa: E402

KEY_MAP = {
    "embed_tokens.weight": "model.embed_tokens.weight",
    "norm.weight": "model.norm.weight",
    "lm_head.weight": "model.lm_head.weight",
}

LAYER_KEY_MAP = {
    "input_layernorm.weight": "input_layernorm.weight",
    "self_attn.q_proj.weight": "self_attn.q_proj.weight",
    "self_attn.k_proj.weight": "self_attn.k_proj.weight",
    "self_attn.v_proj.weight": "self_attn.v_proj.weight",
    "self_attn.o_proj.weight": "self_attn.o_proj.weight",
    "post_attention_layernorm.weight": "post_attention_layernorm.weight",
    "mlp.gate_proj.weight": "mlp.gate_proj.weight",
    "mlp.up_proj.weight": "mlp.up_proj.weight",
    "mlp.down_proj.weight": "mlp.down_proj.weight",
    "timestep_residual.proj.weight": "timestep_residual.proj.weight",
    "timestep_residual.proj.bias": "timestep_residual.proj.bias",
}

TIMESTEP_KEY_MAP = {
    "timestep_emb.mlp.0.weight": "model.timestep_emb.mlp.0.weight",
    "timestep_emb.mlp.0.bias": "model.timestep_emb.mlp.0.bias",
    "timestep_emb.mlp.2.weight": "model.timestep_emb.mlp.2.weight",
    "timestep_emb.mlp.2.bias": "model.timestep_emb.mlp.2.bias",
}

GENERATION_CONFIG = {
    "bos_token_id": 0,
    "eos_token_id": 2,
    "pad_token_id": 1,
    "mask_token_id": 32000,
    "temperature": 0.7,
    "repetition_penalty": 1.5,
    "re_mask": 0.1,
    "num_steps": 128,
    "max_new_tokens": 96,
    "use_cache": False,
    "transformers_version": "4.40.0",
}


def remap_state_dict(state_dict, half=False):
    # Strip torch.compile's _orig_mod. prefix (old checkpoints saved from a
    # compiled model have it; training now saves clean keys)
    state_dict = {k.replace("_orig_mod.", "", 1) if k.startswith("_orig_mod.") else k: v
                  for k, v in state_dict.items()}
    new_dict = {}
    for key, tensor in state_dict.items():
        if key in KEY_MAP:
            new_key = KEY_MAP[key]
        elif key in TIMESTEP_KEY_MAP:
            new_key = TIMESTEP_KEY_MAP[key]
        elif key.startswith("layers."):
            parts = key.split(".")
            # layers.N.<component>.<rest>
            layer_idx, component = parts[1], parts[2]
            rest = ".".join(parts[3:])
            comp_key = f"{component}.{rest}" if rest else component
            new_key = f"model.layers.{layer_idx}.{comp_key}"
        else:
            new_key = key
        new_dict[new_key] = tensor.half() if half else tensor.float()
    return new_dict


def package_scripts(out):
    """Copy the runnable pipeline into out/scripts/ so the release is
    self-contained: chat, finetune, and re-export work from the artifact."""
    src = Path(__file__).resolve().parent
    root = src.parent  # model.py lives in the project root
    scripts_dir = out / "scripts"
    scripts_dir.mkdir(parents=True, exist_ok=True)
    for name in ["model.py", "chat.py", "prepare_data.py", "train_chat.py",
                 "export_hf.py", "convert_data.py"]:
        cand = src / name if (src / name).exists() else root / name
        if cand.exists():
            shutil.copy2(cand, scripts_dir / name)
    req = scripts_dir / "requirements.txt"
    if not req.exists():
        req.write_text("torch>=2.2\ntransformers>=4.40\nsafetensors>=0.4\n"
                       "datasets>=2.18\nnumpy>=1.26\n"
                       "pandas>=2.0\npyarrow>=14.0\n")
    print(f"[*] Packaged scripts -> {scripts_dir}")


def export(checkpoint_path, tokenizer_dir, output_dir, half=False):
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    print(f"[*] Loading checkpoint {checkpoint_path}")
    ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    config = MetaDiffusionConfig(
        **{k: v for k, v in ckpt["config"].items()
           if k in MetaDiffusionConfig.__dataclass_fields__}
    )
    config.tie_word_embeddings = False

    print("[*] Remapping state dict...")
    state_dict = remap_state_dict(ckpt["model_state_dict"], half=half)
    save_file(state_dict, out / "model.safetensors")
    print(f"[*] Saved {len(state_dict)} tensors -> {out / 'model.safetensors'} "
          f"({'fp16' if half else 'fp32'})")

    # Keep vocab fields consistent with the actual weights: the chat pipeline
    # resizes embed_tokens/lm_head to 32010 (Supra 32000 + chat tokens), and
    # the ckpt config still carries vocab_size=32000 from the base. Any loader
    # using vocab_size to size embeddings would fail with a size mismatch.
    n_vocab = state_dict["model.lm_head.weight"].shape[0]
    config.vocab_size = n_vocab
    config.mask_vocab_size = n_vocab
    print(f"[*] Vocab in config: {n_vocab} (matches weights)")

    config_dict = asdict(config)
    config_dict.pop("dtype", None)  # transformers chokes on "torch.float32" strings
    config_dict["model_type"] = "metadiffusion"
    config_dict["architectures"] = ["MetaDiffusionForCausalLM"]
    with open(out / "config.json", "w") as f:
        json.dump(config_dict, f, indent=2)

    with open(out / "generation_config.json", "w") as f:
        json.dump(GENERATION_CONFIG, f, indent=2)

    # Copy tokenizer (has ChatML + rainbow tokens)
    tok_dir = Path(tokenizer_dir)
    for name in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"]:
        src = tok_dir / name
        if src.exists():
            shutil.copy2(src, out / name)

    package_scripts(out)
    print(f"[*] Exported to {out}")
    print("    Vocab:", len(state_dict.get("model.lm_head.weight", [])),
          "| step:", ckpt.get("step"))


def main():
    parser = argparse.ArgumentParser(description="Export chat-SFT checkpoint to HF-style dir")
    parser.add_argument("--checkpoint", required=True, help="step_*.pt or best.pt")
    parser.add_argument("--tokenizer", default="data/no_robots_chatml/tokenizer")
    parser.add_argument("--output", required=True, help="Output dir")
    parser.add_argument("--fp16", action="store_true",
                        help="Save weights as fp16 (half the size; matches the "
                             "base release format)")
    args = parser.parse_args()
    export(args.checkpoint, args.tokenizer, args.output, half=args.fp16)


if __name__ == "__main__":
    main()