from src.pipeline.base import PipelineBase
from typing import Dict, Any, Callable
from queue import Queue, Empty
from src.utils import get_module_logger
from src.models import AudioBinary_data_list
import threading

logger = get_module_logger(__name__)


class ASRPipeline(PipelineBase):
    """
    管道类
    实现具体的处理逻辑
    """

    def __init__(self, *args, **kwargs):
        """
        初始化管道
        """
        super().__init__(*args, **kwargs)
        self._config: Dict[str, Any] = {}
        self._functor_dict: Dict[str, Any] = {}
        self._subqueue_dict: Dict[str, Any] = {}
        self._is_baked: bool = False
        self._input_queue: Queue = None
        self._audio_binary_data_list: AudioBinary_data_list = None

        self._status_lock = threading.Lock()
        self._is_running: bool = False
        self._stop_event: bool = False
        self._callback: Callable = None

    def set_config(self, config: Dict[str, Any]) -> None:
        """
        设置配置
        参数:
            config: Dict[str, Any] 配置
        """
        self._config = config

    def get_config(self) -> Dict[str, Any]:
        """
        获取配置
        返回:
            Dict[str, Any] 配置
        """
        return self._config

    def set_audio_binary(self, audio_binary: AudioBinary_data_list) -> None:
        """
        设置音频二进制存储单元
        参数:
            audio_binary: 音频二进制
        """
        self._audio_binary = audio_binary

    def set_models(self, models: Dict[str, Any]) -> None:
        """
        设置模型
        """
        self._models = models

    def set_input_queue(self, input_queue: Queue) -> None:
        """
        设置输入队列
        """
        self._input_queue = input_queue

    def set_callback(self, callback: Callable) -> None:
        """
        设置回调函数
        """
        self._callback = callback

    def bake(self) -> None:
        """
        烘焙管道
        """
        self._pre_check_resource()
        self._init_functor()
        self._is_baked = True

    def _pre_check_resource(self) -> None:
        """
        预检查资源
        """
        if self._input_queue is None:
            raise RuntimeError("[ASRpipeline]输入队列未设置")
        if self._functor_dict is None:
            raise RuntimeError("[ASRpipeline]functor字典未设置")
        if self._subqueue_dict is None:
            raise RuntimeError("[ASRpipeline]子队列字典未设置")
        if self._audio_binary is None:
            raise RuntimeError("[ASRpipeline]音频存储单元未设置")

    def _init_functor(self) -> None:
        """
        初始化函数
        自身的functor流程图如下
        self.input_queue(self.run检测输入到subqueue["original"])->vad
            ->vad2asr ->asrend
            ->vad2spk ->spkend
                ->asrend+spkend->resultbinder
        ->self.callback     
        """
        try:
            from src.functor import FunctorFactory

            # 加载VAD、asr、spk functor
            self._functor_dict["vad"] = FunctorFactory.make_functor(
                functor_name="vad", config=self._config, models=self._models
            )
            self._functor_dict["asr"] = FunctorFactory.make_functor(
                functor_name="asr", config=self._config, models=self._models
            )
            self._functor_dict["spk"] = FunctorFactory.make_functor(
                functor_name="spk", config=self._config, models=self._models
            )

            self._functor_dict["resultbinder"] = FunctorFactory.make_functor(
                functor_name="resultbinder", config=self._config, models=self._models
            )

            # 创建音频数据存储单元
            self._audio_binary_data_list = AudioBinary_data_list()

            self._functor_dict["vad"].set_audio_binary_data_list(
                self._audio_binary_data_list
            )

            # 初始化子队列
            self._subqueue_dict["original"] = Queue()
            self._subqueue_dict["vad2asr"] = Queue()
            self._subqueue_dict["vad2spk"] = Queue()
            self._subqueue_dict["asrend"] = Queue()
            self._subqueue_dict["spkend"] = Queue()
            # 输出队列
            self._subqueue_dict["OUTPUT"] = Queue()

            # 设置子队列的输入队列
            self._functor_dict["vad"].set_input_queue(self._subqueue_dict["original"])
            self._functor_dict["asr"].set_input_queue(self._subqueue_dict["vad2asr"])
            self._functor_dict["spk"].set_input_queue(self._subqueue_dict["vad2spk"])
            # 设置resultbinder的输入队列
            # 汇总 asr语音识别结果 和 说话人识别结果
            self._functor_dict["resultbinder"].add_input_queue(
                "asr", self._subqueue_dict["asrend"]
            )
            self._functor_dict["resultbinder"].add_input_queue(
                "spk", self._subqueue_dict["spkend"]
            )

            # 设置回调函数——放置到对应队列中
            self._functor_dict["vad"].add_callback(self._subqueue_dict["vad2asr"].put)
            self._functor_dict["vad"].add_callback(self._subqueue_dict["vad2spk"].put)
            # 设置asr与spk的回调函数
            self._functor_dict["asr"].add_callback(self._subqueue_dict["asrend"].put)
            self._functor_dict["spk"].add_callback(self._subqueue_dict["spkend"].put)
            # 设置resultbinder的回调函数 为 自身被设置的回调函数，用于和外界交互
            self._functor_dict["resultbinder"].add_callback(self._callback)

        except ImportError:
            raise ImportError("functorFactory引入失败,ASRPipeline无法完成初始化")

    def _check_result(self, result: Any) -> None:
        """
        检查结果
        """
        # 若asr和spk队列中都有数据，则合并数据
        if (
            self._subqueue_dict["asrend"].qsize()
            & self._subqueue_dict["spkend"].qsize()
        ):
            asr_data = self._subqueue_dict["asrend"].get()
            spk_data = self._subqueue_dict["spkend"].get()
            # 合并数据
            result = {"asr_data": asr_data, "spk_data": spk_data}
            # 通知回调函数
            self._notify_callbacks(result)

    def run(self) -> threading.Thread:
        """
        运行管道
        Returns:
            threading.Thread: 返回已运行线程实例
        """
        # 检查运行资源是否准备完毕
        self._pre_check()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        logger.info("[ASRpipeline]管道开始运行")
        return self._thread

    def _pre_check(self) -> None:
        """
        预检查
        """
        if self._is_baked is False:
            raise RuntimeError("[ASRpipeline]管道未烘焙，无法运行")

        for functor_name, functor in self._functor_dict.items():
            if functor is None:
                raise RuntimeError(f"[ASRpipeline]functor{functor_name}异常")

        for subqueue_name, subqueue in self._subqueue_dict.items():
            if subqueue is None:
                raise RuntimeError(f"[ASRpipeline]子队列{subqueue_name}异常")

    def _run(self) -> None:
        """
        真实的运行逻辑
        """
        # 运行所有functor
        for functor_name, functor in self._functor_dict.items():
            logger.info(f"[ASRpipeline]运行{functor_name}functor")
            self._functor_dict[functor_name].run()

        # 设置管道运行状态
        with self._status_lock:
            self._is_running = True
            self._stop_event = False

        while self._is_running and not self._stop_event:
            try:
                data = self._input_queue.get(timeout=self._queue_timeout)
                # 检查是否是结束信号
                if data is None:
                    logger.info("收到结束信号，管道准备停止")
                    self._input_queue.task_done()  # 标记结束信号已处理
                    break

                # 处理数据
                self._process(data)

                # 标记任务完成
                self._input_queue.task_done()

            except Empty:
                # 队列获取超时，继续等待
                continue
            except Exception as e:
                logger.error(f"[ASRpipeline]管道处理数据出错: {str(e)}")
                break

        logger.info("[ASRpipeline]管道停止运行")

    def _process(self, data: Any) -> Any:
        """
        处理数据
        参数:
            data: 输入数据
        返回:
            处理结果
        """
        # 子类实现具体的处理逻辑
        self._subqueue_dict["original"].put(data)

    def stop(self) -> None:
        """
        停止管道
        """
        with self._status_lock:
            self._is_running = False
            self._stop_event = True
            for functor_name, functor in self._functor_dict.items():
                # logger.info(f"停止{functor_name}functor")
                if functor.stop():
                    logger.info(f"[ASRpipeline]子Functor[{functor_name}]停止")
                else:
                    logger.error(f"[ASRpipeline]子Functor[{functor_name}]停止失败")
            self._thread.join()
            logger.info("[ASRpipeline]管道停止")
            return True
