import argparse
from pathlib import Path

import numpy as np
from rknn.api import RKNN


def parse_args():
    parser = argparse.ArgumentParser(description="Convert Qwen3-ASR audio encoder ONNX to RKNN.")
    parser.add_argument(
        "--onnx-path",
        type=str,
        default="rknn_deploy/audio_encoder/onnx/qwen3_asr_audio_chunk100.onnx",
        help="Path to exported ONNX model.",
    )
    parser.add_argument("--target-platform", type=str, default="rk3588", help="RK target platform.")
    parser.add_argument("--chunk-frames", type=int, default=100, help="Fixed mel chunk length.")
    parser.add_argument(
        "--savepath",
        type=str,
        default="rknn_deploy/audio_encoder/rknn/qwen3_asr_audio_chunk100_rk3588.rknn",
        help="Output RKNN path.",
    )
    return parser.parse_args()


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

    rknn = RKNN(verbose=False)
    rknn.config(target_platform=args.target_platform)

    ret = rknn.load_onnx(
        model=args.onnx_path,
        inputs=["input_features", "feature_len"],
        input_size_list=[[1, 128, args.chunk_frames], [1]],
        input_initial_val=[None, np.asarray([args.chunk_frames], dtype=np.int32)],
    )
    if ret != 0:
        raise SystemExit(f"rknn.load_onnx failed: {ret}")

    ret = rknn.build(do_quantization=False, dataset=None)
    if ret != 0:
        raise SystemExit(f"rknn.build failed: {ret}")

    ret = rknn.export_rknn(savepath.as_posix())
    if ret != 0:
        raise SystemExit(f"rknn.export_rknn failed: {ret}")

    print(f"saved: {savepath}")


if __name__ == "__main__":
    main()