import argparse
import pathlib
import urllib.request

FILES = {
    "whisper_encoder_base_20s.onnx": "https://ftrg.zbox.filez.com/v2/delivery/data/95f00b0fc900458ba134f8b180b3f7a1/examples/whisper/whisper_encoder_base_20s.onnx",
    "whisper_decoder_base_20s.onnx": "https://ftrg.zbox.filez.com/v2/delivery/data/95f00b0fc900458ba134f8b180b3f7a1/examples/whisper/whisper_decoder_base_20s.onnx",
    "mel_80_filters.txt": "https://raw.githubusercontent.com/airockchip/rknn_model_zoo/master/examples/whisper/model/mel_80_filters.txt",
    "vocab_en.txt": "https://raw.githubusercontent.com/airockchip/rknn_model_zoo/master/examples/whisper/model/vocab_en.txt",
    "vocab_zh.txt": "https://raw.githubusercontent.com/airockchip/rknn_model_zoo/master/examples/whisper/model/vocab_zh.txt",
}


def download_file(url: str, dst: pathlib.Path) -> None:
    with urllib.request.urlopen(url, timeout=120) as response:
        data = response.read()
    dst.write_bytes(data)


def main() -> None:
    parser = argparse.ArgumentParser(description="Download Whisper model assets")
    parser.add_argument("--target", default="/models", help="Destination directory")
    parser.add_argument(
        "--force",
        action="store_true",
        help="Re-download files even if they already exist",
    )
    args = parser.parse_args()

    target = pathlib.Path(args.target)
    target.mkdir(parents=True, exist_ok=True)

    for name, url in FILES.items():
        path = target / name
        if path.exists() and not args.force:
            print(f"skip {name} (exists)")
            continue
        print(f"download {name}")
        download_file(url, path)

    print(f"done: {target}")


if __name__ == "__main__":
    main()