#!/usr/bin/env python3
"""Convert a VocalRender checkpoint into a standalone VoxCPM2 GGUF.

Requires torch, safetensors, and an audiocpp_gguf executable built with VoxCPM2.
Example:
    python convert_vocalrender.py --input ./VocalRender-Pro \
        --output ./vocalrender-pro-orig.gguf --converter ./audiocpp_gguf
"""

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

import torch
from safetensors import safe_open
from safetensors.torch import save_file


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True, help="Source checkpoint directory.")
    parser.add_argument("--output", type=Path, required=True, help="Output orig-precision GGUF.")
    parser.add_argument("--converter", default="audiocpp_gguf", help="Converter executable path or PATH command.")
    parser.add_argument("--model-spec", type=Path, help="Optional voxcpm2 spec; otherwise use the converter catalog.")
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    source = args.input.expanduser().resolve(strict=True)
    output = args.output.expanduser().resolve()
    converter = shutil.which(os.path.expanduser(args.converter))
    if converter is None:
        parser.error("Converter not found; provide --converter or add audiocpp_gguf to PATH.")
    if output.exists() and not args.overwrite:
        parser.error("Output already exists; use --overwrite to replace it.")
    weights = source / "model.safetensors"
    with safe_open(weights, framework="pt", device="cpu") as tensors:
        shape = tensors.get_slice("base_lm.embed_tokens.weight").get_shape()
    config = json.loads((source / "config.json").read_text(encoding="utf-8"))
    tokenizer = json.loads((source / "tokenizer.json").read_text(encoding="utf-8"))
    token_config = json.loads((source / "tokenizer_config.json").read_text(encoding="utf-8"))
    ids = [*tokenizer["model"]["vocab"].values(), *[t["id"] for t in tokenizer["added_tokens"]]]
    if len(shape) != 2 or shape[1] != config["lm_config"]["hidden_size"]:
        raise ValueError("Embedding shape does not match the model configuration.")
    if min(ids) < 0 or max(ids) >= shape[0]:
        raise ValueError("Tokenizer IDs exceed the checkpoint embedding capacity.")
    config["lm_config"]["vocab_size"] = shape[0]
    tokenizer["model"]["merges"] = [
        " ".join(merge) if isinstance(merge, list) else merge
        for merge in tokenizer["model"]["merges"]
    ]
    token_config["added_tokens_decoder"] = {
        str(t["id"]): {key: value for key, value in t.items() if key != "id"}
        for t in tokenizer["added_tokens"]
    }
    special = {key: token_config[key] for key in ("bos_token", "eos_token", "unk_token")}
    output.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="vocalrender-convert-", dir=output.parent) as temporary:
        root = Path(temporary) / source.name
        root.mkdir()
        for name, data in (("config.json", config), ("tokenizer.json", tokenizer),
                           ("tokenizer_config.json", token_config), ("special_tokens_map.json", special)):
            (root / name).write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        checkpoint = torch.load(source / "audiovae.pth", map_location="cpu", weights_only=True)
        state = checkpoint.get("state_dict", checkpoint)
        tensors = {}
        for name, tensor in state.items():
            if not isinstance(name, str) or not isinstance(tensor, torch.Tensor):
                raise ValueError("AudioVAE state_dict must contain named tensors only.")
            tensors[name] = tensor.detach().cpu().contiguous()
        if not tensors:
            raise ValueError("AudioVAE checkpoint contains no tensors.")
        vae = root / "audiovae.safetensors"
        save_file(tensors, vae)
        with safe_open(vae, framework="pt", device="cpu") as saved:
            if set(saved.keys()) != set(tensors):
                raise ValueError("AudioVAE tensor names changed during conversion.")
            for name, expected in tensors.items():
                actual = saved.get_tensor(name)
                if actual.dtype != expected.dtype or actual.shape != expected.shape or not torch.equal(actual, expected):
                    raise ValueError(f"AudioVAE tensor changed during conversion: {name}")
        command = [str(Path(converter).resolve()), "--family", "voxcpm2", "--type", "orig",
                   "--input", f"weights={weights}", "--input", f"audiovae_weights={vae}",
                   "--root", str(root), "--output", str(output), "--log"]
        if args.model_spec:
            command.extend(["--model-spec", str(args.model_spec.expanduser().resolve(strict=True))])
        if args.overwrite:
            command.append("--overwrite")
        subprocess.run(command, check=True)
    print(f"Converted: {output}")


if __name__ == "__main__":
    main()