import argparse
from pathlib import Path

import onnxruntime as ort
import torch

from common import StaticChunkAudioEncoder, get_torch_dtype, load_audio_encoder


def parse_args():
    parser = argparse.ArgumentParser(description="Export Qwen3-ASR audio encoder chunk model to ONNX.")
    parser.add_argument("--model-path", type=str, default=".", help="Path to Qwen3-ASR model directory.")
    parser.add_argument(
        "--savepath",
        type=str,
        default="rknn_deploy/audio_encoder/onnx/qwen3_asr_audio_chunk100.onnx",
        help="Output ONNX path.",
    )
    parser.add_argument("--chunk-frames", type=int, default=100, help="Fixed mel chunk length.")
    parser.add_argument(
        "--dtype",
        type=str,
        default="float32",
        choices=["float16", "bfloat16", "float32"],
        help="Torch dtype used for loading and export.",
    )
    parser.add_argument("--device", type=str, default="cpu", help="Torch device for export.")
    return parser.parse_args()


def main():
    args = parse_args()
    savepath = Path(args.savepath)
    savepath.parent.mkdir(parents=True, exist_ok=True)

    tower = load_audio_encoder(model_path=args.model_path, dtype=args.dtype, device=args.device)
    wrapper = StaticChunkAudioEncoder(tower=tower, chunk_frames=args.chunk_frames).to(args.device).eval()

    input_features = torch.zeros(
        (1, 128, args.chunk_frames),
        dtype=get_torch_dtype(args.dtype),
        device=args.device,
    )
    feature_len = torch.tensor([args.chunk_frames], dtype=torch.int32, device=args.device)

    with torch.no_grad():
        torch_features, torch_valid_len = wrapper(input_features, feature_len)

    torch.onnx.export(
        wrapper,
        (input_features, feature_len),
        savepath.as_posix(),
        input_names=["input_features", "feature_len"],
        output_names=["audio_features", "valid_len"],
        opset_version=18,
        dynamo=False,
        dynamic_axes={
            "input_features": {0: "batch"},
            "feature_len": {0: "batch"},
            "audio_features": {0: "batch"},
            "valid_len": {0: "batch"},
        },
    )

    ort_session = ort.InferenceSession(savepath.as_posix(), providers=["CPUExecutionProvider"])
    ort_features, ort_valid_len = ort_session.run(
        None,
        {
            "input_features": input_features.detach().cpu().numpy(),
            "feature_len": feature_len.detach().cpu().numpy(),
        },
    )
    max_diff = float((torch_features.detach().cpu() - torch.from_numpy(ort_features)).abs().max().item())

    print(f"saved: {savepath}")
    print(f"chunk_frames: {args.chunk_frames}")
    print(f"chunk_output_frames: {wrapper.max_aftercnn_len}")
    print(f"torch_valid_len: {int(torch_valid_len[0].item())}")
    print(f"ort_valid_len: {int(ort_valid_len.reshape(-1)[0])}")
    print(f"max_abs_diff(torch_vs_onnx): {max_diff:.8f}")


if __name__ == "__main__":
    main()