"""Build a naive layer-pruned Voxtral draft from Voxtral-Mini-3B-2507.

Loads the full model, keeps a subset of the 30 Llama decoder layers, saves
the pruned model to disk with a corrected config so HF `.from_pretrained()`
reloads it as a smaller `VoxtralForConditionalGeneration`.

Design:
  - Full model: 30 decoder layers in `model.language_model.model.layers`.
  - Keep layers by 0-indexed positions (default: [0, 6, 12, 18, 24, 29]).
  - Preserve everything else unchanged: audio encoder, multi_modal_projector,
    embed_tokens, norm, lm_head. This ensures the draft shares the audio path
    and vocab with the target (both required for HF `assistant_model=`).
  - Save to `--out-dir` so HF can reload with `.from_pretrained()`.

Usage:
    python build_draft.py --keep 0,6,12,18,24,29 --out-dir /mnt/drafts/voxtral-mini-3B-draft-6of30
    python build_draft.py --keep 0,10,20,29 --out-dir /mnt/drafts/voxtral-mini-3B-draft-4of30

Then load in a benchmark script:
    target = VoxtralForConditionalGeneration.from_pretrained("mistralai/Voxtral-Mini-3B-2507", ...)
    draft  = VoxtralForConditionalGeneration.from_pretrained("/mnt/drafts/voxtral-mini-3B-draft-6of30", ...)
    outputs = target.generate(..., assistant_model=draft)
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path

import torch
from transformers import (
    AutoProcessor,
    VoxtralForConditionalGeneration,
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--source", default="mistralai/Voxtral-Mini-3B-2507")
    p.add_argument("--out-dir", required=True, type=Path)
    p.add_argument(
        "--keep",
        default="0,6,12,18,24,29",
        help="Comma-separated 0-indexed layer indices to keep from the 30 "
             "Llama decoder layers. Must be sorted ascending, must include "
             "at least the first and last for stability.",
    )
    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
    p.add_argument("--device", default="cuda")
    return p.parse_args()


def main() -> int:
    args = parse_args()

    keep_layers = [int(x) for x in args.keep.split(",")]
    n_keep = len(keep_layers)
    assert keep_layers == sorted(keep_layers), "--keep must be ascending"
    assert n_keep >= 2, "Keep at least 2 layers"

    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]

    print(f"[build_draft] Loading {args.source} on {args.device} dtype={args.dtype}")
    model = VoxtralForConditionalGeneration.from_pretrained(
        args.source,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
    )
    processor = AutoProcessor.from_pretrained(args.source)

    # Introspect layer container. transformers 5.x path:
    #   VoxtralForConditionalGeneration
    #     .model (VoxtralModel)
    #       .language_model (LlamaModel)
    #         .layers (ModuleList of 30 LlamaDecoderLayer)
    #   .lm_head
    # transformers 4.5x path had .language_model directly under the top-level
    # object (no intermediate .model). Support both.
    if hasattr(model, "model") and hasattr(model.model, "language_model"):
        decoder = model.model.language_model
    elif hasattr(model, "language_model"):
        lm = model.language_model
        decoder = lm.model if hasattr(lm, "model") else lm
    else:
        raise AttributeError(
            "Cannot find decoder. Model has: " + ", ".join(sorted(dict(model.named_children()).keys()))
        )
    assert hasattr(decoder, "layers"), (
        f"Decoder container {type(decoder).__name__} has no .layers. "
        f"Children: {list(dict(decoder.named_children()).keys())}"
    )
    layers = decoder.layers
    n_full = len(layers)
    print(f"[build_draft] Full model has {n_full} decoder layers")
    assert max(keep_layers) < n_full, f"--keep {keep_layers} exceeds n_full={n_full}"

    # Rewrite layer_idx on each kept layer so RoPE / KV cache accounting still works
    kept = torch.nn.ModuleList([layers[i] for i in keep_layers])
    for new_idx, layer in enumerate(kept):
        if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "layer_idx"):
            layer.self_attn.layer_idx = new_idx
    decoder.layers = kept

    # Update text_config so config.json reflects the new depth
    tcfg = model.config.text_config
    tcfg.num_hidden_layers = n_keep
    # Some HF paths look for the num_hidden_layers on the top config too.
    if hasattr(model.config, "num_hidden_layers"):
        model.config.num_hidden_layers = n_keep

    # Save
    args.out_dir.mkdir(parents=True, exist_ok=True)
    print(f"[build_draft] Saving pruned model to {args.out_dir}")
    model.save_pretrained(args.out_dir, safe_serialization=True)
    processor.save_pretrained(args.out_dir)

    # Sanity print
    with open(args.out_dir / "config.json") as f:
        saved = json.load(f)
    print(f"[build_draft] saved text_config.num_hidden_layers = "
          f"{saved.get('text_config', {}).get('num_hidden_layers')}")
    print(f"[build_draft] Done.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())