"""
Export a Wisp checkpoint into a Hugging Face repository layout.

The release is deliberately split in two, because the two halves have very
different compatibility stories.

* **The trunk** is an ordinary Llama decoder: RMSNorm, SwiGLU, grouped query
  attention, RoPE with the half split (neox) convention. Exported under Llama
  parameter names with a `LlamaForCausalLM` config, it loads in `transformers`,
  `mlx_lm`, and anything that converts from those, with no trust_remote_code and
  no custom modelling file. Somebody who just wants a small FIM completion model
  gets one that works in their existing runtime.
* **The MTP module** has no standard home. It ships as a separate
  `mtp.safetensors` plus `mtp_config.json`, which the runtimes that can use it
  can load, and everything else can ignore without breaking.

A model whose weights only load through a bespoke script is a model nobody runs.

Verification is not optional here. Name mapping and RoPE convention are exactly
the kind of silent conversion error that leaves a model quietly dumber rather
than visibly broken, so `--verify` reloads the exported tensors back through the
MLX model and asserts the logits are identical. That catches a wrong mapping. It
does not catch a wrong RoPE convention, which needs a cross framework check
against `transformers`; `scripts/verify_export_torch.py` does that when torch
and transformers are installed.

Usage:
    python scripts/export_hf.py --ckpt out/run1/ckpt_latest \\
        --tokenizer tokenizer/code32k.json --out export/wisp-coder-110m \\
        --repo-id YOUR_NAMESPACE/wisp-coder-110m \\
        --validation-report out/run1/final_validation.json \\
        --acceptance-comparison out/run1/acceptance.control-comparison.json \\
        --format-ablation-report \
          out/run2-no-fim/acceptance.format-ablation.json \\
        --rollout-report out/run1/rollout.registered.v3.json --verify
"""

import argparse
import json
import os
import shutil
import sys
import tempfile

import mlx.core as mx
from mlx.utils import tree_flatten, tree_unflatten

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

from model import Wisp, ModelArgs, causal_mask  # noqa: E402
from scripts.hf_metadata import (  # noqa: E402
    development_evaluation_section,
    file_sha256,
    generation_config,
    llama_config,
    mtp_config,
    render_evaluation_section,
    render_model_card,
    special_tokens_map,
    tokenizer_config,
    validate_export_checkpoint,
    verify_export_manifest,
    write_export_manifest,
    write_json,
)


def load_json_report(path, label):
    if not path:
        raise ValueError(f"final export requires --{label}")
    if not os.path.isfile(path):
        raise FileNotFoundError(f"{label} does not exist: {path}")
    with open(path, encoding="utf-8") as f:
        value = json.load(f)
    if not isinstance(value, dict):
        raise ValueError(f"{label} must contain a JSON object")
    return value, file_sha256(path)


def map_trunk(flat, args):
    """Wisp parameter names to Llama parameter names."""
    out = {}
    out["model.embed_tokens.weight"] = flat["tok_emb.weight"]
    out["model.norm.weight"] = flat["norm.weight"]
    if not args.tie_embeddings:
        out["lm_head.weight"] = flat["lm_head.weight"]

    pairs = [
        ("attn_norm.weight", "input_layernorm.weight"),
        ("attn.wq.weight", "self_attn.q_proj.weight"),
        ("attn.wk.weight", "self_attn.k_proj.weight"),
        ("attn.wv.weight", "self_attn.v_proj.weight"),
        ("attn.wo.weight", "self_attn.o_proj.weight"),
        ("ffn_norm.weight", "post_attention_layernorm.weight"),
        ("ffn.w1.weight", "mlp.gate_proj.weight"),
        ("ffn.w3.weight", "mlp.up_proj.weight"),
        ("ffn.w2.weight", "mlp.down_proj.weight"),
    ]
    for i in range(args.n_layers):
        for ours, theirs in pairs:
            out[f"model.layers.{i}.{theirs}"] = flat[f"blocks.{i}.{ours}"]
    return out


def verify_roundtrip(ckpt_dir, exported, args):
    """Reload the exported tensors into the MLX model and diff the logits."""
    model = Wisp(args)
    master = tree_unflatten(
        list(mx.load(os.path.join(ckpt_dir, "master.safetensors")).items())
    )
    model.update(master)
    model.eval()
    mx.eval(model.parameters())

    inverse = {}
    inverse["tok_emb.weight"] = exported["model.embed_tokens.weight"]
    inverse["norm.weight"] = exported["model.norm.weight"]
    pairs = [
        ("attn_norm.weight", "input_layernorm.weight"),
        ("attn.wq.weight", "self_attn.q_proj.weight"),
        ("attn.wk.weight", "self_attn.k_proj.weight"),
        ("attn.wv.weight", "self_attn.v_proj.weight"),
        ("attn.wo.weight", "self_attn.o_proj.weight"),
        ("ffn_norm.weight", "post_attention_layernorm.weight"),
        ("ffn.w1.weight", "mlp.gate_proj.weight"),
        ("ffn.w3.weight", "mlp.up_proj.weight"),
        ("ffn.w2.weight", "mlp.down_proj.weight"),
    ]
    for i in range(args.n_layers):
        for ours, theirs in pairs:
            inverse[f"blocks.{i}.{ours}"] = exported[f"model.layers.{i}.{theirs}"]

    rebuilt = Wisp(args)
    rebuilt.update(tree_unflatten(list(inverse.items())))
    # The MTP module is not part of the trunk export, so carry it across directly.
    rebuilt.mtp.update(model.mtp.parameters())
    rebuilt.eval()
    mx.eval(rebuilt.parameters())

    probe = mx.array([[1, 2, 3, 4, 5, 6, 7, 8]], dtype=mx.int32)
    mask = causal_mask(probe.shape[1], model.norm.weight.dtype)
    a, _, _ = model(probe, mask)
    b, _, _ = rebuilt(probe, mask)
    mx.eval(a, b)
    delta = float(mx.max(mx.abs(a - b)))
    return delta


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--tokenizer", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--repo-id", required=True)
    ap.add_argument("--model-card", default="MODEL_CARD.md")
    ap.add_argument("--license-file", default="LICENSE")
    ap.add_argument("--validation-report")
    ap.add_argument("--acceptance-comparison")
    ap.add_argument("--format-ablation-report")
    ap.add_argument("--rollout-report")
    ap.add_argument("--verify", action="store_true")
    ap.add_argument(
        "--allow-incomplete",
        action="store_true",
        help="permit a development export before max_steps, never for release",
    )
    cli = ap.parse_args()

    meta_path = os.path.join(cli.ckpt, "meta.json")
    master_path = os.path.join(cli.ckpt, "master.safetensors")
    optimizer_path = os.path.join(cli.ckpt, "optimizer.safetensors")
    with open(meta_path) as f:
        meta = json.load(f)
    release_state = validate_export_checkpoint(meta, cli.allow_incomplete)
    if not os.path.isfile(master_path):
        raise FileNotFoundError(f"checkpoint weights do not exist: {master_path}")
    if not os.path.isfile(optimizer_path):
        raise FileNotFoundError(
            f"checkpoint optimizer state does not exist: {optimizer_path}"
        )
    if not os.path.isfile(cli.tokenizer):
        raise FileNotFoundError(f"tokenizer does not exist: {cli.tokenizer}")
    if not os.path.isfile(cli.model_card):
        raise FileNotFoundError(f"model card does not exist: {cli.model_card}")
    if not os.path.isfile(cli.license_file):
        raise FileNotFoundError(f"license does not exist: {cli.license_file}")
    report_args = {
        "validation": (
            cli.validation_report,
            "validation-report",
        ),
        "acceptance_comparison": (
            cli.acceptance_comparison,
            "acceptance-comparison",
        ),
        "format_ablation": (
            cli.format_ablation_report,
            "format-ablation-report",
        ),
        "rollout": (
            cli.rollout_report,
            "rollout-report",
        ),
    }
    if release_state["complete"]:
        reports = {}
        evaluation_sources = {}
        for key, (path, label) in report_args.items():
            report, digest = load_json_report(path, label)
            reports[key] = report
            evaluation_sources[key] = {"sha256": digest}
        evaluation_markdown = render_evaluation_section(
            reports["validation"],
            reports["acceptance_comparison"],
            reports["format_ablation"],
            reports["rollout"],
        )
    else:
        supplied = [
            label for path, label in report_args.values() if path is not None
        ]
        if supplied:
            raise ValueError(
                "development exports cannot claim final reports: "
                + ", ".join(supplied)
            )
        evaluation_sources = None
        evaluation_markdown = development_evaluation_section(
            release_state["step"],
            release_state["max_steps"],
        )
    model_card_template_sha256 = file_sha256(cli.model_card)
    rendered_card = render_model_card(
        cli.model_card,
        cli.repo_id,
        evaluation_markdown,
    )

    out_dir = os.path.abspath(cli.out)
    if os.path.lexists(out_dir):
        raise FileExistsError(
            f"refusing to merge with or replace existing export: {out_dir}"
        )
    parent = os.path.dirname(out_dir)
    os.makedirs(parent, exist_ok=True)

    checkpoint_hashes = {
        "meta_sha256": file_sha256(meta_path),
        "master_sha256": file_sha256(master_path),
        "optimizer_sha256": file_sha256(optimizer_path),
    }
    args = ModelArgs.from_dict(meta["model_args"])

    master = mx.load(master_path)
    flat = dict(master)
    trunk = map_trunk(flat, args)
    missing = [k for k, v in trunk.items() if v is None]
    if missing:
        raise KeyError(f"unmapped trunk tensors: {missing[:5]}")
    trunk_bf16 = {k: v.astype(mx.bfloat16) for k, v in trunk.items()}
    mtp = {k: v.astype(mx.bfloat16) for k, v in flat.items() if k.startswith("mtp.")}
    n_trunk = sum(v.size for v in trunk.values())
    n_mtp = sum(v.size for v in mtp.values())
    staging = tempfile.mkdtemp(
        prefix=f".{os.path.basename(out_dir)}.staging.", dir=parent
    )
    try:
        mx.save_safetensors(
            os.path.join(staging, "model.safetensors"), trunk_bf16
        )
        mx.save_safetensors(os.path.join(staging, "mtp.safetensors"), mtp)
        write_json(os.path.join(staging, "config.json"), llama_config(args))
        write_json(
            os.path.join(staging, "mtp_config.json"), mtp_config(args, meta)
        )
        write_json(
            os.path.join(staging, "tokenizer_config.json"),
            tokenizer_config(args),
        )
        write_json(
            os.path.join(staging, "special_tokens_map.json"),
            special_tokens_map(),
        )
        write_json(
            os.path.join(staging, "generation_config.json"),
            generation_config(),
        )
        shutil.copy(cli.tokenizer, os.path.join(staging, "tokenizer.json"))
        with open(
            os.path.join(staging, "README.md"), "w", encoding="utf-8"
        ) as f:
            f.write(rendered_card)
        shutil.copy(cli.license_file, os.path.join(staging, "LICENSE"))

        write_export_manifest(
            staging,
            meta,
            n_trunk,
            n_mtp,
            cli.repo_id,
            checkpoint_hashes,
            release_state["complete"],
            evaluation_sources,
            model_card_template_sha256,
        )
        verify_export_manifest(staging)
        if (
            file_sha256(meta_path) != checkpoint_hashes["meta_sha256"]
            or file_sha256(master_path) != checkpoint_hashes["master_sha256"]
            or (
                file_sha256(optimizer_path)
                != checkpoint_hashes["optimizer_sha256"]
            )
        ):
            raise RuntimeError("source checkpoint changed during export")

        if cli.verify:
            delta = verify_roundtrip(cli.ckpt, trunk, args)
            status = "identical" if delta == 0.0 else f"max abs delta {delta:.3e}"
            print(f"roundtrip through the exported names: {status}")
            if delta != 0.0:
                raise RuntimeError(
                    "export verification failed, the name mapping is wrong"
                )
            print(
                "NOTE: this proves the name mapping, not the RoPE convention. "
                "Diff against transformers on the same prompt before publishing."
            )
        if os.path.lexists(out_dir):
            raise FileExistsError(
                f"export target appeared during staging: {out_dir}"
            )
        os.rename(staging, out_dir)
        staging = None
    finally:
        if staging is not None and os.path.isdir(staging):
            shutil.rmtree(staging)

    label = "final" if release_state["complete"] else "development snapshot"
    print(
        f"trunk: {len(trunk)} tensors, {n_trunk/1e6:.1f}M params "
        "-> model.safetensors"
    )
    print(
        f"mtp:   {len(mtp)} tensors, {n_mtp/1e6:.1f}M params "
        "-> mtp.safetensors"
    )
    print(f"atomic {label} package plus export_manifest.json -> {out_dir}")


if __name__ == "__main__":
    main()