"""
-*- encoding: utf-8 -*-

ASRRunner
继承RunnerBase
专属pipeline为ASRPipeline
"""

from src.pipeline.ASRpipeline import ASRPipeline
from src.pipeline import PipelineFactory
from src.models import AudioBinary_data_list, AudioBinary_Config
from src.core.model_loader import ModelLoader
from queue import Queue
import soundfile
import time

from src.utils.logger import get_module_logger

logger = get_module_logger(__name__)


OVAERWATCH = False

model_loader = ModelLoader()


def test_asr_pipeline():
    # 加载模型
    args = {
        "asr_model": "paraformer-zh",
        "asr_model_revision": "v2.0.4",
        "vad_model": "fsmn-vad",
        "vad_model_revision": "v2.0.4",
        "spk_model": "cam++",
        "spk_model_revision": "v2.0.2",
        "audio_update": False,
    }
    models = model_loader.load_models(args)
    audio_data, sample_rate = soundfile.read("tests/vad_example.wav")
    audio_config = AudioBinary_Config(
        chunk_size=200,
        chunk_stride=1600,
        sample_rate=sample_rate,
        sample_width=16,
        channels=1,
    )
    chunk_stride = int(audio_config.chunk_size * sample_rate / 1000)
    audio_config.chunk_stride = chunk_stride

    # 创建参数Dict
    config = {
        "audio_config": audio_config,
    }

    # 创建音频数据列表
    audio_binary_data_list = AudioBinary_data_list()

    input_queue = Queue()

    # 创建Pipeline
    # asr_pipeline = ASRPipeline()
    # asr_pipeline.set_models(models)
    # asr_pipeline.set_config(config)
    # asr_pipeline.set_audio_binary(audio_binary_data_list)
    # asr_pipeline.set_input_queue(input_queue)
    # asr_pipeline.add_callback(lambda x: print(f"pipeline callback: {x}"))
    # asr_pipeline.bake()
    asr_pipeline = PipelineFactory.create_pipeline(
        pipeline_name = "ASRpipeline", 
        models=models, 
        config=config, 
        audio_binary=audio_binary_data_list, 
        input_queue=input_queue, 
        callback=lambda x: print(f"pipeline callback: {x}")
    )
    
    # 运行Pipeline
    asr_instance = asr_pipeline.run()


    audio_clip_len = 200
    print(
        f"audio_data: {len(audio_data)}, audio_clip_len: {audio_clip_len}, clip_num: {len(audio_data) // audio_clip_len}"
    )
    for i in range(0, len(audio_data), audio_clip_len):
        input_queue.put(audio_data[i : i + audio_clip_len])

    # time.sleep(10)
    # input_queue.put(None)

    # 等待Pipeline结束
    # asr_instance.join()

    time.sleep(5)
    asr_pipeline.stop()
    # asr_pipeline.stop()



class STTRunner(RunnerBase):
    """
    运行器类
    负责管理资源和协调Pipeline的运行
    """

    def __init__(
        self,
        *args,
        **kwargs,
    ):
        """
        """
        # ws资源
        self._ws_pool: Dict[str,List[WebSocketClient]] = {}
        # 接收资源
        self._audio_binary_list = audio_binary_list
        self._models = models
        self._pipeline_list = pipeline_list

        # 线程控制
        self._lock = Lock()
        # 停止控制
        self._stop_timeout = 10.0
        self._is_stopping = False

        # 配置资源
        for pipeline in self._pipeline_list:
            # 设置输入队列
            pipeline.set_input_queue(self._input_queue)

            # 配置资源
            pipeline.set_audio_binary(
                self._audio_binary_list[pipeline.get_config("audio_binary_name")]
            )
            pipeline.set_models(self._models)

    def run(self) -> None:
        """
        启动所有管道
        """
        logger.info("[%s] 启动所有管道", self.__class__.__name__)
        if not self._pipeline_list:
            raise RuntimeError("没有可用的管道")

        # 启动所有管道
        for pipeline in self._pipeline_list:
            thread = Thread(target=pipeline.run, name=f"Pipeline-{id(pipeline)}")
            thread.daemon = True
            thread.start()
            logger.info("[%s] 管道 %s 已启动", self.__class__.__name__, id(pipeline))

    def stop(self, force: bool = False) -> bool:
        """
        停止所有管道
        参数:
            force: 是否强制停止
        返回:
            bool: 是否成功停止
        """
        if self._is_stopping:
            logger.warning("运行器已经在停止中")
            return False

        self._is_stopping = True
        logger.info("正在停止运行器...")

        try:
            # 发送结束信号
            self._input_queue.put(None)

            # 停止所有管道
            success = True
            for pipeline in self._pipeline_list:
                if force:
                    pipeline.force_stop()
                else:
                    if not pipeline.stop(timeout=self._stop_timeout):
                        logger.warning("管道 %s 停止超时", id(pipeline))
                        success = False

            # 等待队列处理完成
            try:
                start_time = time.time()
                while not self._input_queue.empty():
                    if time.time() - start_time > self._stop_timeout:
                        logger.warning(
                            "等待队列处理完成超时（%s秒）, 队列中还有 %d 个任务未处理",
                            self._stop_timeout,
                            self._input_queue.qsize(),
                        )
                        success = False
                        break
                    time.sleep(0.1)  # 避免过度消耗CPU
            except Exception as e:
                error_type = type(e).__name__
                error_msg = str(e)
                error_traceback = traceback.format_exc()
                logger.error(
                    "等待队列处理完成时发生错误:\n"
                    "错误类型: %s\n"
                    "错误信息: %s\n"
                    "错误堆栈:\n%s",
                    error_type,
                    error_msg,
                    error_traceback,
                )
                success = False

            if success:
                logger.info("所有管道已成功停止")
            else:
                logger.warning(
                    "部分管道停止失败, 队列状态: 大小=%d, 是否为空=%s",
                    self._input_queue.qsize(),
                    self._input_queue.empty(),
                )

            return success

        finally:
            self._is_stopping = False

    def __del__(self) -> None:
        """
        析构函数
        """
        self.stop(force=True)
