"""
VADFunctor
负责对音频片段进行VAD处理, 以VAD_Result进行callback
"""

import threading
from queue import Empty, Queue
from typing import List, Any, Callable
import numpy
from src.models import (
    VAD_Functor_result,
    AudioBinary_Config,
    AudioBinary_data_list,
)
from src.functor.base import BaseFunctor

# 日志
from src.utils.logger import get_module_logger

logger = get_module_logger(__name__)


class VADFunctor(BaseFunctor):
    """
    VADFunctor
    负责对音频片段进行VAD处理, 以VAD_Result进行callback
    需要配置好 _model, _callback, _input_queue, _audio_config, _audio_binary_data_list
    否则无法run()启动线程

    运行中, 使用reset_cache()重置缓存, 准备下次任务

    使用stop()停止线程, 但需要等待input_queue为空
    """

    def __init__(self) -> None:
        super().__init__()
        # 资源与配置
        self._model: dict = {}  # 模型
        self._callback: List[Callable] = []  # 回调函数
        self._input_queue: Queue = None  # 输入队列
        self._audio_config: AudioBinary_Config = None  # 音频配置
        self._audio_binary_data_list: AudioBinary_data_list = None  # 音频数据列表

        # flag
        # 此处用到两个锁，但都是为了截断_run线程，考虑后续优化
        self._is_running: bool = False
        self._stop_event: bool = False

        # 线程资源
        self._thread: threading.Thread = None

        # 状态锁
        self._status_lock: threading.Lock = threading.Lock()

        # 缓存
        self._audio_cache: numpy.ndarray = None
        self._audio_cache_preindex: int = 0
        self._model_cache: dict = {}
        self._cache_result_list = []
        self._audiobinary_cache = None

    def reset_cache(self) -> None:
        """
        重置缓存, 用于任务完成后清理缓存数据, 准备下次任务
        """
        self._audio_cache = None
        self._audio_cache_preindex = 0
        self._model_cache = {}
        self._cache_result_list = []
        self._audiobinary_cache = None

    def set_input_queue(self, queue: Queue) -> None:
        """
        设置监听的输入消息队列
        """
        self._input_queue = queue

    def set_model(self, model: dict) -> None:
        """
        设置推理模型
        """
        self._model = model

    def set_audio_config(self, audio_config: AudioBinary_Config) -> None:
        """
        设置音频配置
        """
        self._audio_config = audio_config
        logger.debug("VADFunctor设置音频配置: %s", self._audio_config)

    def set_audio_binary_data_list(
        self, audio_binary_data_list: AudioBinary_data_list
    ) -> None:
        """
        设置音频数据列表, 为Class AudioBinary_data_list类型
        AudioBinary_data_list包含binary_data_list, 为list[_AudioBinary_data]类型
        _AudioBinary_data包含binary_data, 为bytes/numpy.ndarray类型
        """
        self._audio_binary_data_list = audio_binary_data_list

    def add_callback(self, callback: Callable) -> None:
        """
        向自身的_callback: List[Callable]回调函数列表中添加回调函数
        """
        if not isinstance(self._callback, list):
            self._callback = []
        self._callback.append(callback)

    def _do_callback(self, result: List[List[int]]) -> None:
        """
        回调函数
        VADFunctor包装结果, 存储到AudioBinary中, 并向队列中添加AudioBinary_Slice

        输入:
            result: List[[start,end]] 处理所得VAD端点
            其中若start==-1, 则表示前无端点, 若end==-1, 则表示后无端点
            当处理得到一个完成片段时, 存入AudioBinary中, 并向队列中添加AudioBinary_Slice
        输出:
            None
        """
        # 持久化缓存结果队列
        for pair in result:
            [start, end] = pair
            # 若无前端点, 则向缓存队列中合并
            if start == -1:
                self._cache_result_list[-1][1] = end
            else:
                self._cache_result_list.append(pair)
        while len(self._cache_result_list) > 1:
            # 创建VAD片段
            # 计算开始帧
            start_frame = self._audio_config.ms2frame(self._cache_result_list[0][0])
            start_frame -= self._audio_cache_preindex
            # 计算结束帧
            end_frame = self._audio_config.ms2frame(self._cache_result_list[0][1])
            end_frame -= self._audio_cache_preindex
            # 计算开始时间
            vad_result = VAD_Functor_result.create_from_push_data(
                audiobinary_data_list=self._audio_binary_data_list,
                data=self._audiobinary_cache[start_frame:end_frame],
                start_time=self._cache_result_list[0][0],
                end_time=self._cache_result_list[0][1],
            )
            self._audio_cache_preindex += end_frame
            self._audiobinary_cache = self._audiobinary_cache[end_frame:]
            for callback in self._callback:
                callback(vad_result)
            self._cache_result_list.pop(0)

    def _predeal_data(self, data: Any) -> None:
        """
        预处理数据, 将数据缓存到_audio_cache和_audiobinary_cache中
        """
        if self._audio_cache is None:
            self._audio_cache = data
        else:
            # 拼接音频数据
            if isinstance(self._audio_cache, numpy.ndarray):
                self._audio_cache = numpy.concatenate((self._audio_cache, data))
            elif isinstance(self._audio_cache, list):
                self._audio_cache.append(data)
        if self._audiobinary_cache is None:
            self._audiobinary_cache = data
        else:
            # 拼接音频数据
            if isinstance(self._audiobinary_cache, numpy.ndarray):
                self._audiobinary_cache = numpy.concatenate(
                    (self._audiobinary_cache, data)
                )
            elif isinstance(self._audiobinary_cache, list):
                self._audiobinary_cache.append(data)

    def _process(self, data: Any):
        """
        处理数据
        使用model进行生成, 并使用_do_callback进行回调
        """
        self._predeal_data(data)
        if len(self._audio_cache) >= self._audio_config.chunk_stride:
            result = self._model["vad"].generate(
                input=self._audio_cache,
                cache=self._model_cache,
                chunk_size=self._audio_config.chunk_size,
                is_final=False,
            )
            if len(result[0]["value"]) > 0:
                self._do_callback(result[0]["value"])
                # logger.debug(f"VADFunctor结果: {result[0]['value']}")
            self._audio_cache = None

    def _run(self):
        """
        线程运行逻辑
        监听输入队列, 当有数据时, 处理数据
        当输入队列为空时, 间隔1s检测是否进入停止事件。
        """
        # 刷新运行状态
        with self._status_lock:
            self._is_running = True
            self._stop_event = False
        # 运行逻辑
        while self._is_running:
            try:
                data = self._input_queue.get(True, timeout=1)
                self._process(data)
                self._input_queue.task_done()
            # 当队列为空时, 间隔1s检测是否进入停止事件。
            except Empty:
                if self._stop_event:
                    break
                continue
            # 其他异常
            except Exception as e:
                logger.error("VADFunctor运行时发生错误: %s", e)
                raise e

    def run(self):
        """
        启动 _run 线程, 并返回线程对象
        """
        self._pre_check()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return self._thread

    def _pre_check(self) -> bool:
        """
        检测硬性资源是否都已设置
        """
        if self._model is None:
            raise ValueError("模型未设置")
        if self._audio_config is None:
            raise ValueError("音频配置未设置")
        if self._audio_binary_data_list is None:
            raise ValueError("音频数据列表未设置")
        if self._input_queue is None:
            raise ValueError("输入队列未设置")
        if self._callback is None:
            raise ValueError("回调函数未设置")
        return True

    def stop(self):
        """
        停止线程
        通过设置_stop_event为True, 来在input_queue.get()循环为空时退出
        """
        with self._status_lock:
            self._stop_event = True
        self._thread.join()
        with self._status_lock:
            self._is_running = False
        return not self._thread.is_alive()


# class VAD:

#     def __init__(
#         self,
#         VAD_model=None,
#         audio_config: AudioBinary_Config = None,
#         callback: Callable = None,
#     ):
#         # vad model
#         self.VAD_model = VAD_model
#         if self.VAD_model is None:
#             self.VAD_model = AutoModel(
#                 model="fsmn-vad", model_revision="v2.0.4", disable_update=True
#             )
#         # audio config
#         self.audio_config = audio_config
#         # vad result
#         self.vad_result = VADResponse(time_chunk_index_callback=callback)
#         # audio binary poll
#         self.audio_chunk = AudioChunk(audio_config=self.audio_config)
#         self.cache = {}

#     def push_binary_data(
#         self,
#         binary_data: bytes,
#     ):
#         # 压入二进制数据
#         self.audio_chunk.add_chunk(binary_data)
#         # 处理音频块
#         res = self.VAD_model.generate(
#             input=binary_data,
#             cache=self.cache,
#             chunk_size=self.audio_config.chunk_size,
#             is_final=False,
#         )
#         # print("VAD generate", res)
#         if len(res[0]["value"]):
#             self.vad_result += VADResponse.from_raw(res)

#     def set_callback(
#         self,
#         callback: Callable,
#     ):
#         self.vad_result.time_chunk_index_callback = callback

#     def process_vad_result(self, callback: Callable = None):
#         # 处理VAD结果
#         callback = (
#             callback
#             if callback is not None
#             else self.vad_result.time_chunk_index_callback
#         )
#         self.vad_result.process_time_chunk(
#             lambda x: callback(
#                 AudioBinary_Chunk(
#                     start_time=x["start_time"],
#                     end_time=x["end_time"],
#                     chunk=self.audio_chunk.get_chunk(x["start_time"], x["end_time"]),
#                 )
#             )
#         )
