#!/usr/bin/env python3
"""Export mistralai/Voxtral-Mini-3B-2507 to the onnx-asr "speech-llm" three-graph contract.

Graphs
------
encoder.onnx      input_features (1,128,3000*N) -> audio_embeds (1,375*N,3072)
embed_tokens.onnx input_ids (1,S) -> inputs_embeds (1,S,3072)
decoder.onnx      inputs_embeds (1,S,H) + attn_bias (1,1,S,P+S) + position_ids (1,S)
                  + past_key_values.{i}.{key,value} (1,KV,P,D)
                  -> logits (1,S,V) + present.{i}.{key,value} (1,KV,P+S,D)

The Voxtral audio encoder is a Whisper encoder: full attention over a fixed
3000-frame (30 s) mel window, so the graph needs neither packing indices nor a
block-diagonal bias and the runtime feeds the features unchanged.
"""

from __future__ import annotations

import argparse
import base64
import json
import shutil
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn

REPO = "mistralai/Voxtral-Mini-3B-2507"

# The eight languages Voxtral-Mini-3B-2507 is documented to transcribe.
LANGUAGES = {
    "en": "English",
    "fr": "French",
    "de": "German",
    "es": "Spanish",
    "it": "Italian",
    "pt": "Portuguese",
    "nl": "Dutch",
    "hi": "Hindi",
}

MEL_FRAMES = 3000


class EncoderExport(nn.Module):
    """Whisper-style audio encoder + Voxtral multi-modal projector.

    The mel features arrive as one flat window of `3000 * N` frames. They are
    split into the 30 s chunks the encoder expects, encoded, and the resulting
    1500 frames per chunk are grouped four at a time (the projector input width
    is four times the encoder width) before the projection into the language
    model embedding space.
    """

    def __init__(self, encoder: nn.Module, projector: nn.Module, num_mel_bins: int, intermediate_size: int):
        super().__init__()
        self.encoder = encoder
        self.projector = projector
        self.num_mel_bins = num_mel_bins
        self.intermediate_size = intermediate_size

    def forward(self, input_features: torch.Tensor) -> torch.Tensor:
        enc = self.encoder
        chunks = (
            input_features.reshape(1, self.num_mel_bins, -1, MEL_FRAMES)
            .permute(0, 2, 1, 3)
            .reshape(-1, self.num_mel_bins, MEL_FRAMES)
        )

        hidden = F.gelu(enc.conv1(chunks))
        hidden = F.gelu(enc.conv2(hidden))
        hidden = hidden.permute(0, 2, 1)
        hidden = hidden + enc.embed_positions.weight

        for layer in enc.layers:
            hidden = layer(hidden, attention_mask=None)

        hidden = enc.layer_norm(hidden)
        hidden = hidden.reshape(-1, self.intermediate_size)
        return self.projector(hidden).unsqueeze(0)


class EmbedExport(nn.Module):
    def __init__(self, embed_tokens: nn.Module):
        super().__init__()
        self.embed_tokens = embed_tokens

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.embed_tokens(input_ids)


class DecoderExport(nn.Module):
    """Llama causal LM over inputs_embeds with a flat past/present KV cache."""

    def __init__(self, language_model: nn.Module, lm_head: nn.Module, num_layers: int):
        super().__init__()
        self.language_model = language_model
        self.lm_head = lm_head
        self.num_layers = num_layers

    def forward(self, inputs_embeds, attn_bias, position_ids, *past):
        from transformers.cache_utils import DynamicCache

        cache = DynamicCache(config=self.language_model.config)
        for i in range(self.num_layers):
            cache.update(past[2 * i], past[2 * i + 1], i)

        out = self.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attn_bias,
            position_ids=position_ids,
            past_key_values=cache,
            use_cache=True,
        )
        logits = self.lm_head(out.last_hidden_state)
        present: list[torch.Tensor] = []
        for i in range(self.num_layers):
            present.append(cache.layers[i].keys)
            present.append(cache.layers[i].values)
        return (logits, *present)


MAX_PROTOBUF = 1_900_000_000


def consolidate(src: Path, dst: Path) -> None:
    """Move a graph to `dst`, spilling weights to `<name>.onnx_data` when it is too big for protobuf."""
    import onnx

    raw_size = sum(path.stat().st_size for path in src.parent.rglob("*") if path.is_file())
    model = onnx.load(str(src), load_external_data=True)
    dst.parent.mkdir(parents=True, exist_ok=True)
    data_path = dst.with_suffix(".onnx_data")
    if data_path.exists():
        data_path.unlink()
    if raw_size > MAX_PROTOBUF:
        onnx.save(
            model,
            str(dst),
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=data_path.name,
            size_threshold=1024,
            convert_attribute=False,
        )
    else:
        onnx.save(model, str(dst))


def export_graph(module, inputs, name: str, out: Path, tmp: Path, **kwargs) -> None:
    stage = tmp / name
    stage.mkdir(parents=True, exist_ok=True)
    torch.onnx.export(module, inputs, str(stage / f"{name}.onnx"), do_constant_folding=True, dynamo=False, **kwargs)
    consolidate(stage / f"{name}.onnx", out / f"{name}.onnx")
    shutil.rmtree(stage, ignore_errors=True)


def quantize(name: str, out: Path, tmp: Path) -> None:
    from onnxruntime.quantization import QuantType, quantize_dynamic

    stage = tmp / f"{name}_int8"
    stage.mkdir(parents=True, exist_ok=True)
    # Graphs whose fp32 weights already live in a side file stay over the 2 GB
    # protobuf limit after quantization of the non-MatMul tensors, so onnxruntime
    # has to write external data too.
    quantize_dynamic(
        out / f"{name}.onnx",
        stage / f"{name}_int8.onnx",
        weight_type=QuantType.QInt8,
        use_external_data_format=(out / f"{name}.onnx_data").exists(),
        extra_options={"MatMulConstBOnly": True},
    )
    consolidate(stage / f"{name}_int8.onnx", out / f"{name}_int8.onnx")
    shutil.rmtree(stage, ignore_errors=True)


def bytes_to_unicode() -> dict[int, str]:
    """GPT-2 byte -> printable unicode table (same one onnx-asr uses to decode)."""
    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("\xa1"), ord("\xac") + 1)) + list(
        range(ord("\xae"), ord("\xff") + 1)
    )
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    return dict(zip(bs, [chr(c) for c in cs]))


def build_vocab(repo: str, vocab_size: int) -> dict[str, int]:
    """Build a {token: id} map straight from tekken.json.

    `MistralCommonBackend.get_vocab()` decodes every piece as UTF-8, which is
    lossy for the byte-level pieces, so the raw token bytes are read from
    tekken.json and written in the GPT-2 byte-to-unicode encoding that the
    onnx-asr byte-level decoder reverses exactly.
    """
    from huggingface_hub import hf_hub_download

    with Path(hf_hub_download(repo, "tekken.json")).open("rt", encoding="utf-8") as f:
        tekken = json.load(f)

    num_special = tekken["config"]["default_num_special_tokens"]
    byte_encoder = bytes_to_unicode()

    vocab: dict[str, int] = {}
    for special in tekken["special_tokens"][:num_special]:
        vocab[special["token_str"]] = special["rank"]
    for entry in tekken["vocab"][: vocab_size - num_special]:
        token = "".join(byte_encoder[b] for b in base64.b64decode(entry["token_bytes"]))
        vocab[token] = entry["rank"] + num_special

    assert len(vocab) == vocab_size, (len(vocab), vocab_size)
    return vocab


def build_prompts(processor, repo: str, audio_token_id: int, clip: Path) -> tuple[list[int], list[int], dict]:
    """Read the transcription-request token ids straight out of mistral-common.

    Voxtral encodes a transcription request as
    `<s> [INST] <begin_audio> <audio>*375 <end_audio> [lang tokens] [/INST]`,
    so the language marker sits in the suffix, after the audio.
    """

    def encode(language: str | None) -> list[int]:
        out = processor.apply_transcription_request(
            audio=str(clip), model_id=repo, language=language, return_tensors="pt", tokenize=True, return_dict=True
        )
        return out["input_ids"][0].tolist()

    def split(ids: list[int]) -> tuple[list[int], list[int]]:
        first = ids.index(audio_token_id)
        last = len(ids) - 1 - ids[::-1].index(audio_token_id)
        assert last - first + 1 == ids.count(audio_token_id), "audio tokens are not contiguous"
        return ids[:first], ids[last + 1 :]

    prefix, suffix = split(encode(None))
    suffixes: dict[str, list[int]] = {}
    for code, name in LANGUAGES.items():
        lang_prefix, lang_suffix = split(encode(code))
        assert lang_prefix == prefix, (code, lang_prefix, prefix)
        suffixes[code] = lang_suffix
        suffixes[name] = lang_suffix
    return prefix, suffix, suffixes


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path, default=Path(__file__).parent)
    parser.add_argument("--repo", default=REPO)
    parser.add_argument("--opset", type=int, default=17)
    parser.add_argument("--only", nargs="*", default=None)
    parser.add_argument("--skip-int8", action="store_true")
    parser.add_argument("--skip-export", action="store_true", help="only (re)run quantization")
    args = parser.parse_args()

    from transformers import AutoConfig, AutoProcessor, VoxtralForConditionalGeneration

    out = args.out
    out.mkdir(parents=True, exist_ok=True)
    tmp = out / "_staging"
    wanted = set(args.only) if args.only else {"encoder", "embed_tokens", "decoder", "config"}
    if args.skip_export:
        wanted -= {"config"}

    config = AutoConfig.from_pretrained(args.repo)
    text_config = config.text_config
    audio_config = config.audio_config
    num_layers = text_config.num_hidden_layers
    head_dim = text_config.head_dim
    num_kv = text_config.num_key_value_heads
    hidden = text_config.hidden_size

    processor = AutoProcessor.from_pretrained(args.repo)

    if wanted - {"config"} and not args.skip_export:
        model = VoxtralForConditionalGeneration.from_pretrained(
            args.repo, dtype=torch.float32, attn_implementation="sdpa"
        ).eval()
        inner = model.model

    if "encoder" in wanted and not args.skip_export:
        print("exporting encoder...", flush=True)
        enc_wrapper = EncoderExport(
            inner.audio_tower, inner.multi_modal_projector, audio_config.num_mel_bins, audio_config.intermediate_size
        ).eval()
        export_graph(
            enc_wrapper,
            (torch.randn(1, audio_config.num_mel_bins, MEL_FRAMES),),
            "encoder",
            out,
            tmp,
            input_names=["input_features"],
            output_names=["audio_embeds"],
            dynamic_axes={"input_features": {2: "padded_frames"}, "audio_embeds": {1: "audio_seq"}},
            opset_version=args.opset,
        )

    if "embed_tokens" in wanted and not args.skip_export:
        print("exporting embed_tokens...", flush=True)
        export_graph(
            EmbedExport(inner.language_model.embed_tokens).eval(),
            (torch.zeros(1, 5, dtype=torch.int64),),
            "embed_tokens",
            out,
            tmp,
            input_names=["input_ids"],
            output_names=["inputs_embeds"],
            dynamic_axes={"input_ids": {1: "seq"}, "inputs_embeds": {1: "seq"}},
            opset_version=args.opset,
        )

    if "decoder" in wanted and not args.skip_export:
        print("exporting decoder...", flush=True)
        dec_wrapper = DecoderExport(inner.language_model, model.lm_head, num_layers).eval()
        s, p = 3, 2
        dummy_past = []
        for _ in range(num_layers):
            dummy_past.append(torch.randn(1, num_kv, p, head_dim))
            dummy_past.append(torch.randn(1, num_kv, p, head_dim))

        past_names, present_names = [], []
        for i in range(num_layers):
            past_names += [f"past_key_values.{i}.key", f"past_key_values.{i}.value"]
            present_names += [f"present.{i}.key", f"present.{i}.value"]

        dyn = {
            "inputs_embeds": {1: "seq"},
            "attn_bias": {2: "seq", 3: "total"},
            "position_ids": {1: "seq"},
            "logits": {1: "seq"},
        }
        for name in past_names:
            dyn[name] = {2: "past"}
        for name in present_names:
            dyn[name] = {2: "total"}

        export_graph(
            dec_wrapper,
            (
                torch.randn(1, s, hidden),
                torch.zeros(1, 1, s, p + s),
                torch.arange(p, p + s, dtype=torch.int64)[None],
                *dummy_past,
            ),
            "decoder",
            out,
            tmp,
            input_names=["inputs_embeds", "attn_bias", "position_ids", *past_names],
            output_names=["logits", *present_names],
            dynamic_axes=dyn,
            opset_version=args.opset,
        )

    if "config" in wanted and not args.skip_export:
        print("writing vocab and config...", flush=True)
        vocab = build_vocab(args.repo, text_config.vocab_size)
        with (out / "vocab.json").open("wt", encoding="utf-8") as f:
            json.dump(vocab, f, ensure_ascii=False)

        prefix_ids, suffix_ids, language_suffixes = build_prompts(
            processor, args.repo, config.audio_token_id, out / "clips" / "en_1.wav"
        )
        asr_config = {
            "model_type": "speech-llm",
            "features_size": audio_config.num_mel_bins,
            "preprocessor": f"whisper{audio_config.num_mel_bins}",
            "hidden_size": hidden,
            "num_layers": num_layers,
            "num_key_value_heads": num_kv,
            "head_dim": head_dim,
            "eos_token_ids": [config.text_config.eos_token_id or 2],
            "max_sequence_length": 512,
            "prompt_prefix_ids": prefix_ids,
            "prompt_suffix_ids": suffix_ids,
            "language_suffix_ids": language_suffixes,
            "tokenizer_type": "byte-level",
            "source_model": args.repo,
        }
        with (out / "config.json").open("wt", encoding="utf-8") as f:
            json.dump(asr_config, f, ensure_ascii=False, indent=2)

    if not args.skip_int8:
        for name in ("encoder", "embed_tokens", "decoder"):
            if name in wanted:
                print(f"quantizing {name}...", flush=True)
                quantize(name, out, tmp)

    shutil.rmtree(tmp, ignore_errors=True)

    for path in sorted(out.glob("*.onnx*")):
        print(path.name, path.stat().st_size)


if __name__ == "__main__":
    main()