"""
Functor测试
VAD测试
"""
from src.functor.vad_functor import VADFunctor
from queue import Queue, Empty
from src.model_loader import ModelLoader
from src.models import AudioBinary_Config
from src.utils.data_format import wav_to_bytes
import time
from src.utils.logger import get_module_logger

logger = get_module_logger(__name__)

model_loader = ModelLoader()

def test_vad_functor():
    # 加载模型
    args = {
        "vad_model": "fsmn-vad",
        "vad_model_revision": "v2.0.4",
        "auto_update": False,
    }
    model_loader.load_models(args)
    # 创建VAD函数器
    vad_functor = VADFunctor()
    # 创建输入队列
    input_queue = Queue()
    # 设置输入队列
    vad_functor.set_input_queue(input_queue)
    # 设置音频配置
    vad_functor.set_audio_config(AudioBinary_Config(
        chunk_size=960,
        chunk_stride=480,
        sample_rate=16000,
        sample_width=2,
        channels=1
    ))
    # 设置回调函数
    vad_functor.add_callback(lambda x: print(x))
    # 设置模型
    vad_functor.set_model({
        'vad': model_loader.models['vad']
    })
    # 启动VAD函数器
    vad_functor.run()

    # 加载数据
    f_binary = wav_to_bytes("tests/vad_example.wav")
    chunk_size = 960000
    # chunk_size = len(f_binary)
    print(f"f_binary: {len(f_binary)}, chunk_size: {chunk_size}, clip_num: {len(f_binary) // chunk_size}")
    for i in range(0, len(f_binary), chunk_size):
        binary_data = f_binary[i:i+chunk_size]
        input_queue.put(binary_data)
    # 等待VAD函数器结束
    time.sleep(10)
    vad_functor.stop()