import argparse
from StreamVAD import StreamVAD
import os
import librosa
import soundfile as sf


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str, required=True, help='Input audio file')
    parser.add_argument('--backend', type=str, default='ax650', choices=['ax650', 'ax630c', 'onnx'])
    parser.add_argument('--output_dir', type=str, default='output', help='Output audio dir')
    return parser.parse_args()


def main():
    args = get_args()
    os.makedirs(args.output_dir, exist_ok=True)

    vad = StreamVAD(args.backend, 
                    sensitivity=0.5,
                    silence_ms=200)
    
    audio, _ = librosa.load(args.input, sr=vad.model.sr, mono=True)
    i = 0
    for result in vad.run(audio, vad.model.sr):
        if result:
            print(result)
            filename = os.path.join(args.output_dir, f"{i}.wav")
            sf.write(filename, result['audio'], samplerate=vad.model.sr)
            i += 1


if __name__ == '__main__':
    main()