# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import gc
import os
import re
import tarfile

import torch
import yaml
from tokenizers import AddedToken

from transformers import (
    ParakeetCTCConfig,
    ParakeetEncoderConfig,
    ParakeetFeatureExtractor,
    ParakeetForCTC,
    ParakeetForRNNT,
    ParakeetForTDT,
    ParakeetProcessor,
    ParakeetRNNTConfig,
    ParakeetTDTConfig,
    ParakeetTokenizer,
)
from transformers.convert_slow_tokenizer import ParakeetConverter
from transformers.utils.hub import cached_file


NEMO_TO_HF_WEIGHT_MAPPING = {
    r"encoder\.pre_encode\.conv\.": r"encoder.subsampling.layers.",
    r"encoder\.pre_encode\.out\.": r"encoder.subsampling.linear.",
    r"encoder\.pos_enc\.": r"encoder.encode_positions.",
    r"encoder\.layers\.(\d+)\.conv\.batch_norm\.": r"encoder.layers.\1.conv.norm.",
    r"decoder\.decoder_layers\.0\.(weight|bias)": r"ctc_head.\1",
    r"linear_([kv])": r"\1_proj",
    r"linear_out": r"o_proj",
    r"linear_q": r"q_proj",
    r"pos_bias_([uv])": r"bias_\1",
    r"linear_pos": r"relative_k_proj",
}

# Additional mappings for TDT decoder and joint network
NEMO_TDT_WEIGHT_MAPPING = {
    r"decoder\.prediction\.embed\.": r"decoder.embedding.",
    r"decoder\.prediction\.dec_rnn\.lstm\.": r"decoder.lstm.",
    r"joint\.enc\.": r"encoder_projector.",
    r"joint\.pred\.": r"decoder.decoder_projector.",
    r"joint\.joint_net\.2\.": r"joint.head.",
}


def convert_key(key, mapping):
    for pattern, replacement in mapping.items():
        key = re.sub(pattern, replacement, key)
    return key


def extract_nemo_archive(nemo_file_path: str, extract_dir: str) -> dict[str, str]:
    """Extract .nemo file (tar archive) and return paths to important files."""
    print(f"Extracting NeMo archive: {nemo_file_path}")

    with tarfile.open(nemo_file_path, "r", encoding="utf-8") as tar:
        tar.extractall(extract_dir)

    all_files = []
    for root, dirs, files in os.walk(extract_dir):
        for file in files:
            file_path = os.path.join(root, file)
            all_files.append(file_path)

    print(f"All extracted files: {[os.path.basename(f) for f in all_files]}")

    model_files = {}
    for root, dirs, files in os.walk(extract_dir):
        for file in files:
            file_path = os.path.join(root, file)
            file_lower = file.lower()

            if (
                file.endswith(".pt")
                or file.endswith(".pth")
                or file.endswith(".ckpt")
                or file.endswith(".bin")
                or "model" in file_lower
                and ("weight" in file_lower or "state" in file_lower)
                or file_lower == "model.pt"
                or file_lower == "pytorch_model.bin"
                or file_lower == "model_weights.ckpt"
            ):
                model_files["model_weights"] = file_path
                print(f"Found model weights: {file}")

            elif (
                file == "model_config.yaml"
                or file == "config.yaml"
                or (file.endswith(".yaml") and "config" in file_lower)
            ):
                if "model_config" not in model_files:
                    model_files["model_config"] = file_path
                    print(f"Found config file: {file}")
                if file == "model_config.yaml":
                    model_files["model_config"] = file_path

            elif (
                file.endswith(".vocab")
                or file.endswith(".model")
                or file.endswith(".txt")
                or ("tokenizer" in file_lower and (file.endswith(".vocab") or file.endswith(".model")))
            ):
                if "tokenizer_model_file" not in model_files or file.endswith(".model"):
                    model_files["tokenizer_model_file"] = file_path
                    print(f"Found tokenizer model file: {file}")
                else:
                    print(f"Found additional vocabulary file (using existing): {file}")

    print(f"Found model files: {list(model_files.keys())}")

    if "model_weights" not in model_files:
        raise FileNotFoundError(
            f"Could not find model weights file in {nemo_file_path}. "
            f"Expected files with extensions: .pt, .pth, .ckpt, .bin. "
            f"Found files: {[os.path.basename(f) for f in all_files]}"
        )

    if "model_config" not in model_files:
        raise FileNotFoundError(
            f"Could not find model config file in {nemo_file_path}. "
            f"Expected: model_config.yaml or config.yaml. "
            f"Found files: {[os.path.basename(f) for f in all_files]}"
        )

    return model_files


def write_processor(
    nemo_config: dict, model_files, output_dir, model_type, push_to_repo_id=None, create_pr=True, revision=None
):
    tokenizer_converted = ParakeetConverter(model_files["tokenizer_model_file"]).converted()
    tokenizer_converted_fast = ParakeetTokenizer(
        tokenizer_object=tokenizer_converted,
        clean_up_tokenization_spaces=False,
    )

    if tokenizer_converted_fast.convert_tokens_to_ids("<unk>") is None:
        # Normally CTC and TDT already have
        tokenizer_converted_fast.add_tokens([AddedToken("<unk>", normalized=False, special=True)])
        print(f"Added <unk> token at ID: {tokenizer_converted_fast.convert_tokens_to_ids('<unk>')}")

    if model_type == "rnnt":
        # RNN-T (unlike TDT) has no dedicated pad token in its NeMo vocab. NeMo's blank is the final vocab entry,
        # i.e. `config.blank_token_id == len(labels)`, and the joint head emits exactly `len(labels) + 1` logits.
        # Add `<blank>` *first* so it lands on that id (no `<pad>` is appended to push it past the model's vocab),
        # and reuse it as the pad token: decoding already skips the pad id, and padding decoder/label tensors with
        # the blank id keeps every id within the joint vocab. This aligns the tokenizer's `<blank>` id with the
        # model's blank logit, so the processor can prepend `<blank>` to build decoder_input_ids unchanged.
        tokenizer_converted_fast.add_tokens([AddedToken("<blank>", normalized=False, special=True)])
        print(f"Added <blank> token at ID: {tokenizer_converted_fast.convert_tokens_to_ids('<blank>')}")
        pad_token = AddedToken("<blank>", normalized=False, special=True)
    else:
        if tokenizer_converted_fast.convert_tokens_to_ids("<pad>") is None:
            # Normally CTC doesn't have while TDT has at token id = 2
            tokenizer_converted_fast.add_tokens([AddedToken("<pad>", normalized=False, special=True)])
            print(f"Added <pad> token at ID: {tokenizer_converted_fast.convert_tokens_to_ids('<pad>')}")
        if model_type == "tdt":
            # TDT needs a separate blank token (its vocab already carries a distinct <pad>)
            tokenizer_converted_fast.add_tokens([AddedToken("<blank>", normalized=False, special=True)])
            print(f"Added <blank> token at ID: {tokenizer_converted_fast.convert_tokens_to_ids('<blank>')}")
        pad_token = AddedToken("<pad>", normalized=False, special=True)
    tokenizer_converted_fast.add_special_tokens(
        {
            "pad_token": pad_token,
            "unk_token": AddedToken("<unk>", normalized=False, special=True),
        }
    )

    feature_extractor_keys_to_ignore = ["_target_", "pad_to", "frame_splicing", "dither", "normalize", "window", "log"]
    feature_extractor_config_keys_mapping = {
        "sample_rate": "sampling_rate",
        "window_size": "win_length",
        "window_stride": "hop_length",
        "window": "window",
        "n_fft": "n_fft",
        "log": "log",
        "features": "feature_size",
        "dither": "dither",
        "pad_to": "pad_to",
        "pad_value": "padding_value",
        "frame_splicing": "frame_splicing",
        "preemphasis": "preemphasis",
        "hop_length": "hop_length",
    }
    converted_feature_extractor_config = {}

    for key, value in nemo_config["preprocessor"].items():
        if key in feature_extractor_keys_to_ignore:
            continue
        if key in feature_extractor_config_keys_mapping:
            if key in ["window_size", "window_stride"]:
                value = int(value * nemo_config["preprocessor"]["sample_rate"])
            converted_feature_extractor_config[feature_extractor_config_keys_mapping[key]] = value
        else:
            raise ValueError(f"Key {key} not found in feature_extractor_keys_mapping")

    feature_extractor = ParakeetFeatureExtractor(**converted_feature_extractor_config)
    processor = ParakeetProcessor(
        feature_extractor=feature_extractor,
        tokenizer=tokenizer_converted_fast,
        # Drives the decoding/timestamp rules (CTC collapses repeats; TDT uses predicted durations + punctuation
        # attach; RNN-T uses single-frame spans).
        decoder_type=model_type,
    )
    processor.save_pretrained(output_dir)

    if push_to_repo_id:
        commit_info = processor.push_to_hub(push_to_repo_id, create_pr=create_pr, revision=revision)
        if create_pr and hasattr(commit_info, "pr_url") and commit_info.pr_url:
            pr_num = commit_info.pr_url.rstrip("/").split("/")[-1]
            return f"refs/pr/{pr_num}"

    return revision


def convert_encoder_config(nemo_config):
    """Convert NeMo encoder config to HF encoder config."""
    encoder_keys_to_ignore = [
        "att_context_size",
        "causal_downsampling",
        "stochastic_depth_start_layer",
        "feat_out",
        "stochastic_depth_drop_prob",
        "_target_",
        "ff_expansion_factor",
        "untie_biases",
        "att_context_style",
        "self_attention_model",
        "conv_norm_type",
        "subsampling",
        "stochastic_depth_mode",
        "conv_context_size",
        "dropout_pre_encoder",
        "reduction",
        "reduction_factor",
        "reduction_position",
    ]
    encoder_config_keys_mapping = {
        "d_model": "hidden_size",
        "n_heads": "num_attention_heads",
        "n_layers": "num_hidden_layers",
        "feat_in": "num_mel_bins",
        "conv_kernel_size": "conv_kernel_size",
        "subsampling_factor": "subsampling_factor",
        "subsampling_conv_channels": "subsampling_conv_channels",
        "pos_emb_max_len": "max_position_embeddings",
        "dropout": "dropout",
        "dropout_emb": "dropout_positions",
        "dropout_att": "attention_dropout",
        "xscaling": "scale_input",
        "use_bias": "attention_bias",
    }
    converted_encoder_config = {}

    for key, value in nemo_config["encoder"].items():
        if key in encoder_keys_to_ignore:
            continue
        if key in encoder_config_keys_mapping:
            converted_encoder_config[encoder_config_keys_mapping[key]] = value
            if key == "use_bias":
                converted_encoder_config["convolution_bias"] = value
        else:
            raise ValueError(f"Key {key} not found in encoder_config_keys_mapping")

    return ParakeetEncoderConfig(**converted_encoder_config)


def load_and_convert_state_dict(model_files):
    """Load NeMo state dict and convert keys to HF format."""
    state_dict = torch.load(model_files["model_weights"], map_location="cpu", weights_only=True)
    converted_state_dict = {}
    for key, value in state_dict.items():
        if key.endswith("featurizer.window") or key.endswith("featurizer.fb"):
            print(f"Skipping preprocessing weight: {key}")
            continue
        converted_key = convert_key(key, NEMO_TO_HF_WEIGHT_MAPPING)
        converted_state_dict[converted_key] = value

    return converted_state_dict


def write_ctc_model(encoder_config, converted_state_dict, output_dir, push_to_repo_id=None, revision=None):
    """Write CTC model using encoder config and converted state dict."""
    model_config = ParakeetCTCConfig(encoder_config=encoder_config)

    print("Loading the checkpoint in a Parakeet CTC model.")
    with torch.device("meta"):
        model = ParakeetForCTC(model_config)
    model.load_state_dict(converted_state_dict, strict=True, assign=True)
    print("Checkpoint loaded successfully.")
    del model.config._name_or_path

    print("Saving the model.")
    model.save_pretrained(output_dir)

    if push_to_repo_id:
        model.push_to_hub(push_to_repo_id, revision=revision)

    del model

    gc.collect()
    print("Reloading the model to check if it's saved correctly.")
    ParakeetForCTC.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
    print("Model reloaded successfully.")


def convert_tdt_config(nemo_config, encoder_config):
    """Convert NeMo TDT config to HF TDT config."""
    decoder_config = nemo_config["decoder"]
    decoding_config = nemo_config["decoding"]
    labels = nemo_config["labels"]
    blank_token_id = len(labels)
    vocab_size = len(labels) + 1  # +1 for blank token, which is added to tokenizer

    prednet = decoder_config.get("prednet", {})
    decoder_hidden_size = prednet.get("pred_hidden", 640)
    num_decoder_layers = prednet.get("pred_rnn_layers", 2)
    durations = decoding_config.get("durations", [0, 1, 2, 3, 4])
    print(
        f"TDT config: vocab_size={vocab_size} (including blank token), "
        f"decoder_hidden={decoder_hidden_size}, "
        f"decoder_layers={num_decoder_layers}, durations={durations}, "
    )

    return ParakeetTDTConfig(
        vocab_size=vocab_size,
        decoder_hidden_size=decoder_hidden_size,
        num_decoder_layers=num_decoder_layers,
        durations=durations,
        hidden_act="relu",
        max_symbols_per_step=10,
        encoder_config=encoder_config.to_dict(),
        pad_token_id=labels.index("<pad>"),
        blank_token_id=blank_token_id,  # blank token is different from pad token for TDT
    )


def load_and_convert_tdt_state_dict(model_files, vocab_size):
    """Load NeMo TDT state dict and convert keys to HF format, splitting combined head."""
    state_dict = torch.load(model_files["model_weights"], map_location="cpu", weights_only=True)
    converted_state_dict = {}

    all_mappings = {**NEMO_TO_HF_WEIGHT_MAPPING, **NEMO_TDT_WEIGHT_MAPPING}

    for key, value in state_dict.items():
        if key.endswith("featurizer.window") or key.endswith("featurizer.fb"):
            print(f"Skipping preprocessing weight: {key}")
            continue

        converted_key = convert_key(key, all_mappings)
        converted_state_dict[converted_key] = value

    return converted_state_dict


def write_tdt_model(nemo_config, encoder_config, model_files, output_dir, push_to_repo_id=None, revision=None):
    """Write TDT model using encoder config, TDT config, and converted state dict."""
    model_config = convert_tdt_config(nemo_config, encoder_config)
    print(f"Converted TDT config: {model_config}")

    converted_state_dict = load_and_convert_tdt_state_dict(model_files, model_config.vocab_size)

    print("Loading the checkpoint in a Parakeet TDT model.")
    with torch.device("meta"):
        model = ParakeetForTDT(model_config)

    missing_keys, unexpected_keys = model.load_state_dict(converted_state_dict, strict=False, assign=True)

    if missing_keys:
        print(f"Warning: Missing keys: {missing_keys}")
    if unexpected_keys:
        print(f"Warning: Unexpected keys: {unexpected_keys}")

    if not missing_keys and not unexpected_keys:
        print("All weights loaded successfully!")

    del model.config._name_or_path

    model.generation_config.decoder_start_token_id = model.config.blank_token_id
    model.generation_config.suppress_tokens = list(
        range(model.config.vocab_size, model.config.vocab_size + len(model.config.durations))
    )

    print("Saving the model.")
    model.save_pretrained(output_dir)

    if push_to_repo_id:
        model.push_to_hub(push_to_repo_id, revision=revision)

    del model

    gc.collect()
    print("Reloading the model to check if it's saved correctly.")
    ParakeetForTDT.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
    print("Model reloaded successfully.")


def convert_rnnt_config(nemo_config, encoder_config):
    """Convert NeMo RNN-T config to HF RNN-T config (a TDT config with no token durations)."""
    decoder_config = nemo_config["decoder"]
    labels = nemo_config["labels"]
    blank_token_id = len(labels)
    vocab_size = len(labels) + 1  # +1 for blank token, which is added to tokenizer

    prednet = decoder_config.get("prednet", {})
    decoder_hidden_size = prednet.get("pred_hidden", 640)
    num_decoder_layers = prednet.get("pred_rnn_layers", 2)
    print(
        f"RNN-T config: vocab_size={vocab_size} (including blank token), "
        f"decoder_hidden={decoder_hidden_size}, decoder_layers={num_decoder_layers}"
    )

    return ParakeetRNNTConfig(
        vocab_size=vocab_size,
        decoder_hidden_size=decoder_hidden_size,
        num_decoder_layers=num_decoder_layers,
        hidden_act="relu",
        max_symbols_per_step=10,
        encoder_config=encoder_config.to_dict(),
        pad_token_id=labels.index("<pad>") if "<pad>" in labels else 0,
        blank_token_id=blank_token_id,  # blank token is different from pad token for RNN-T
    )


def write_rnnt_model(nemo_config, encoder_config, model_files, output_dir, push_to_repo_id=None, revision=None):
    """Write RNN-T model using encoder config, RNN-T config, and converted state dict.

    RNN-T shares the TDT decoder/joint weight layout (the joint head is just `vocab_size` wide, with no
    duration logits), so the TDT state-dict converter is reused as-is.
    """
    model_config = convert_rnnt_config(nemo_config, encoder_config)
    print(f"Converted RNN-T config: {model_config}")

    converted_state_dict = load_and_convert_tdt_state_dict(model_files, model_config.vocab_size)

    print("Loading the checkpoint in a Parakeet RNN-T model.")
    with torch.device("meta"):
        model = ParakeetForRNNT(model_config)

    missing_keys, unexpected_keys = model.load_state_dict(converted_state_dict, strict=False, assign=True)

    if missing_keys:
        print(f"Warning: Missing keys: {missing_keys}")
    if unexpected_keys:
        print(f"Warning: Unexpected keys: {unexpected_keys}")

    if not missing_keys and not unexpected_keys:
        print("All weights loaded successfully!")

    del model.config._name_or_path

    model.generation_config.decoder_start_token_id = model.config.blank_token_id

    print("Saving the model.")
    model.save_pretrained(output_dir)

    if push_to_repo_id:
        model.push_to_hub(push_to_repo_id, revision=revision)

    del model

    gc.collect()
    print("Reloading the model to check if it's saved correctly.")
    ParakeetForRNNT.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
    print("Model reloaded successfully.")


def write_model(nemo_config, model_files, model_type, output_dir, push_to_repo_id=None, revision=None):
    """Main model conversion function."""
    encoder_config = convert_encoder_config(nemo_config)
    print(f"Converted encoder config: {encoder_config}")

    if model_type == "ctc":
        converted_state_dict = load_and_convert_state_dict(model_files)
        write_ctc_model(encoder_config, converted_state_dict, output_dir, push_to_repo_id, revision)
    elif model_type == "tdt":
        write_tdt_model(nemo_config, encoder_config, model_files, output_dir, push_to_repo_id, revision)
    elif model_type == "rnnt":
        write_rnnt_model(nemo_config, encoder_config, model_files, output_dir, push_to_repo_id, revision)
    else:
        raise ValueError(f"Model type {model_type} not supported.")


def main(
    hf_repo_id,
    output_dir,
    model_type,
    push_to_repo_id=None,
    create_pr=True,
    revision=None,
):
    nemo_filename = f"{hf_repo_id.split('/')[-1]}.nemo"
    filepath = cached_file(hf_repo_id, nemo_filename)

    model_files = extract_nemo_archive(filepath, os.path.dirname(filepath))
    nemo_config = yaml.load(open(model_files["model_config"], "r"), Loader=yaml.FullLoader)

    # When revision is given (e.g. "refs/pr/3"), both pushes target that existing PR branch.
    # Otherwise, write_processor creates a new PR and returns its revision for write_model.
    pr_revision = write_processor(
        nemo_config,
        model_files,
        output_dir,
        model_type,
        push_to_repo_id,
        create_pr=create_pr if revision is None else False,
        revision=revision,
    )
    write_model(nemo_config, model_files, model_type, output_dir, push_to_repo_id, pr_revision)


"""
CTC conversion example:
```bash
python src/transformers/models/parakeet/convert_nemo_to_hf.py \
    --hf_repo_id nvidia/parakeet-ctc-1.1b \
    --model_type ctc \
    --output_dir OUTPUT_DIR  \
    --push_to_repo_id USERNAME/parakeet-ctc-1.1b
```

TDT conversion example:
```bash
python src/transformers/models/parakeet/convert_nemo_to_hf.py \
    --hf_repo_id nvidia/parakeet-tdt-0.6b-v3 \
    --model_type tdt \
    --output_dir OUTPUT_DIR  \
    --push_to_repo_id USERNAME/parakeet-tdt-0.6b-v3-hf
```
"""
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--hf_repo_id", required=True, help="Model repo on huggingface.co")
    parser.add_argument(
        "--model_type", required=True, choices=["ctc", "tdt", "rnnt"], help="Model type (`ctc`, `tdt`, `rnnt`)"
    )
    parser.add_argument("--output_dir", required=True, help="Output directory for HuggingFace model")
    parser.add_argument("--push_to_repo_id", help="Repository ID to push the model to on the Hub")
    parser.add_argument(
        "--create_pr",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="Create a PR when pushing to the Hub (default: True). Use --no-create_pr to push directly.",
    )
    parser.add_argument(
        "--revision",
        default=None,
        help='Push to an existing Hub PR branch (e.g. "refs/pr/3"). Overrides --create_pr.',
    )
    args = parser.parse_args()
    main(
        args.hf_repo_id,
        args.output_dir,
        args.model_type,
        args.push_to_repo_id,
        args.create_pr,
        args.revision,
    )
