"""Export parakeet-unified-en-0.6b to ONNX (encoder + decoder_joint).

Downloads the .nemo checkpoint from HuggingFace, exports fp32 ONNX with
consolidated external weights, then produces int8 dynamic-quantized variants.

Outputs:
  onnx_fp32/
    encoder.onnx         + encoder.onnx.data         (~2.4 GB)
    decoder_joint.onnx   + decoder_joint.onnx.data   (~34 MB)
    tokenizer.model

  onnx_int8/
    encoder.int8.onnx                                 (~624 MB)
    decoder_joint.int8.onnx                           (~9 MB)
    tokenizer.model

Requirements:
    pip install -r requirements.txt
"""
import os
import shutil
import tempfile

import onnx
from huggingface_hub import hf_hub_download
import nemo.collections.asr as nemo_asr
from omegaconf import OmegaConf, open_dict
from onnxruntime.quantization import QuantType, quantize_dynamic

HF_REPO = "nvidia/parakeet-unified-en-0.6b"
NEMO_FILENAME = "parakeet-unified-en-0.6b.nemo"
OPSET = 17
FP32_DIR = "onnx_fp32"
INT8_DIR = "onnx_int8"


# --- download .nemo -----------------------------------------------------------
print(f"downloading {NEMO_FILENAME} from {HF_REPO}…")
nemo_path = hf_hub_download(repo_id=HF_REPO, filename=NEMO_FILENAME)
print(f"  {nemo_path}")

# --- load model --------------------------------------------------------------
print("loading model…")
m = nemo_asr.models.ASRModel.restore_from(nemo_path, map_location="cpu")
with open_dict(m.cfg):
    if m.cfg.get("validation_ds") is None:
        m.cfg.validation_ds = OmegaConf.create({})
m.eval()

os.makedirs(FP32_DIR, exist_ok=True)
os.makedirs(INT8_DIR, exist_ok=True)


# --- fp32 export -------------------------------------------------------------
def export_fp32(module, name: str) -> str:
    out_path = os.path.join(FP32_DIR, name)
    with tempfile.TemporaryDirectory() as tmp:
        staged = os.path.join(tmp, name)
        module.export(staged, onnx_opset_version=OPSET)
        proto = onnx.load(staged, load_external_data=True)
        onnx.save(
            proto,
            out_path,
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=f"{name}.data",
            size_threshold=1024,
            convert_attribute=False,
        )
    print(f"  {out_path} (+ {name}.data)")
    return out_path


# --- int8 quantization -------------------------------------------------------
def quantize_int8(fp32_path: str, name: str) -> None:
    int8_name = name.replace(".onnx", ".int8.onnx")
    out_path = os.path.join(INT8_DIR, int8_name)
    with tempfile.TemporaryDirectory() as tmp:
        proto = onnx.load(fp32_path, load_external_data=True)
        # ensure ai.onnx domain is explicit (ort quantizer requires it)
        for op in proto.opset_import:
            if op.domain == "":
                op.domain = "ai.onnx"
        # save to tmp with external data (encoder exceeds 2GB protobuf limit)
        staged = os.path.join(tmp, name)
        onnx.save(
            proto,
            staged,
            save_as_external_data=True,
            all_tensors_to_one_file=True,
            location=f"{name}.data",
            size_threshold=1024,
            convert_attribute=False,
        )
        quantize_dynamic(
            model_input=staged,
            model_output=out_path,
            weight_type=QuantType.QInt8,
        )
    print(f"  {out_path}")


# --- extract tokenizer -------------------------------------------------------
def extract_tokenizer() -> None:
    with tempfile.TemporaryDirectory() as tmp:
        import tarfile

        with tarfile.open(nemo_path, "r") as tar:
            for member in tar.getmembers():
                if member.name.endswith("_tokenizer.model"):
                    tar.extract(member, tmp)
                    src = os.path.join(tmp, member.name)
                    for d in (FP32_DIR, INT8_DIR):
                        shutil.copy2(src, os.path.join(d, "tokenizer.model"))
                    print(f"  tokenizer.model -> {FP32_DIR}/, {INT8_DIR}/")
                    return
    print("  WARNING: tokenizer.model not found in .nemo archive")


# --- run ---------------------------------------------------------------------
print("exporting fp32…")
enc_fp32 = export_fp32(m.encoder, "encoder.onnx")
dj_fp32 = export_fp32(m.decoder_joint, "decoder_joint.onnx")

print("quantizing int8…")
quantize_int8(enc_fp32, "encoder.onnx")
quantize_int8(dj_fp32, "decoder_joint.onnx")

print("extracting tokenizer…")
extract_tokenizer()

print("done.")