import argparse
import os
from whisper_ax import Whisper
import time


def get_args():
    parser = argparse.ArgumentParser(
        prog="whisper", description="Run Whisper on input audio file"
    )
    parser.add_argument("--wav", "-w", type=str, required=True, help="Input audio file")
    parser.add_argument(
        "--model_type",
        "-t",
        type=str,
        choices=["tiny", "base", "small", "large", "large-v3", "turbo"],
        required=True,
        help="model type, only support tiny, base and small currently",
    )
    parser.add_argument(
        "--model_path",
        "-p",
        type=str,
        required=False,
        default="../models-ax650",
        help="model path for *.axmodel, tokens.txt",
    )
    parser.add_argument(
        "--language",
        "-l",
        type=str,
        required=False,
        default="zh",
        help="Target language, support en, zh, ja, and others. See languages.py for more options.",
    )
    parser.add_argument(
        "--task",
        type=str,
        required=False,
        choices=["translate", "transcribe"],
        default="transcribe",
    )
    return parser.parse_args()


def main():
    args = get_args()
    print(vars(args))

    # Check wav existence
    wav_path = args.wav
    assert os.path.exists(wav_path), f"{wav_path} NOT exist"

    model = Whisper(args.model_type, args.model_path, args.language, args.task)

    print("ASR result:")
    start = time.time()
    print(model.run(wav_path))
    end = time.time()

    import librosa

    samples, sr = librosa.load(wav_path, sr=16000)
    duration = len(samples) / sr
    process_time = end - start
    print(f"RTF: {process_time / duration}")


if __name__ == "__main__":
    main()