"""Convert RVC WebUI's fairseq HuBERT checkpoint to Transformers format."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from transformers import HubertConfig, HubertModel


VOICECHANGER_ROOT = Path(__file__).resolve().parents[1]


def resolve_path(path: str) -> Path:
    value = Path(path)
    if value.is_absolute():
        return value
    return VOICECHANGER_ROOT / value


def map_state_dict(fairseq_state: dict[str, torch.Tensor], hf_state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    mapped: dict[str, torch.Tensor] = {}

    def put(hf_key: str, fs_key: str) -> None:
        value = fairseq_state[fs_key]
        if tuple(value.shape) != tuple(hf_state[hf_key].shape):
            raise ValueError(
                f"Shape mismatch: {fs_key} {tuple(value.shape)} -> "
                f"{hf_key} {tuple(hf_state[hf_key].shape)}"
            )
        mapped[hf_key] = value

    if "masked_spec_embed" in hf_state and "mask_emb" in fairseq_state:
        put("masked_spec_embed", "mask_emb")

    for idx in range(7):
        put(f"feature_extractor.conv_layers.{idx}.conv.weight", f"feature_extractor.conv_layers.{idx}.0.weight")
    put("feature_extractor.conv_layers.0.layer_norm.weight", "feature_extractor.conv_layers.0.2.weight")
    put("feature_extractor.conv_layers.0.layer_norm.bias", "feature_extractor.conv_layers.0.2.bias")

    put("feature_projection.layer_norm.weight", "layer_norm.weight")
    put("feature_projection.layer_norm.bias", "layer_norm.bias")
    put("feature_projection.projection.weight", "post_extract_proj.weight")
    put("feature_projection.projection.bias", "post_extract_proj.bias")

    put("encoder.pos_conv_embed.conv.bias", "encoder.pos_conv.0.bias")
    put("encoder.pos_conv_embed.conv.parametrizations.weight.original0", "encoder.pos_conv.0.weight_g")
    put("encoder.pos_conv_embed.conv.parametrizations.weight.original1", "encoder.pos_conv.0.weight_v")
    put("encoder.layer_norm.weight", "encoder.layer_norm.weight")
    put("encoder.layer_norm.bias", "encoder.layer_norm.bias")

    for idx in range(12):
        prefix_hf = f"encoder.layers.{idx}"
        prefix_fs = f"encoder.layers.{idx}"
        for proj in ("k", "v", "q", "out"):
            put(f"{prefix_hf}.attention.{proj}_proj.weight", f"{prefix_fs}.self_attn.{proj}_proj.weight")
            put(f"{prefix_hf}.attention.{proj}_proj.bias", f"{prefix_fs}.self_attn.{proj}_proj.bias")
        put(f"{prefix_hf}.layer_norm.weight", f"{prefix_fs}.self_attn_layer_norm.weight")
        put(f"{prefix_hf}.layer_norm.bias", f"{prefix_fs}.self_attn_layer_norm.bias")
        put(f"{prefix_hf}.feed_forward.intermediate_dense.weight", f"{prefix_fs}.fc1.weight")
        put(f"{prefix_hf}.feed_forward.intermediate_dense.bias", f"{prefix_fs}.fc1.bias")
        put(f"{prefix_hf}.feed_forward.output_dense.weight", f"{prefix_fs}.fc2.weight")
        put(f"{prefix_hf}.feed_forward.output_dense.bias", f"{prefix_fs}.fc2.bias")
        put(f"{prefix_hf}.final_layer_norm.weight", f"{prefix_fs}.final_layer_norm.weight")
        put(f"{prefix_hf}.final_layer_norm.bias", f"{prefix_fs}.final_layer_norm.bias")

    missing = sorted(set(hf_state) - set(mapped))
    if missing:
        raise ValueError(f"Unmapped Transformers keys: {missing}")
    return mapped


def convert(checkpoint_path: Path, output_dir: Path) -> None:
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    fairseq_state = checkpoint["model"]
    config = HubertConfig(
        hidden_dropout=0.0,
        activation_dropout=0.0,
        attention_dropout=0.0,
        feat_proj_dropout=0.0,
        final_dropout=0.0,
        mask_time_prob=0.0,
        layerdrop=0.0,
    )
    model = HubertModel(config)
    mapped = map_state_dict(fairseq_state, model.state_dict())
    model.load_state_dict(mapped, strict=True)
    model.eval()
    output_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output_dir)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="assets/hubert/hubert_base.pt")
    parser.add_argument("--output", default="assets/hubert/hubert_base_transformers")
    args = parser.parse_args()
    convert(resolve_path(args.input), resolve_path(args.output))
    print(f"saved Transformers HuBERT to {args.output}")


if __name__ == "__main__":
    main()