"""
语音活动检测模块 - 基于 SileroVAD
检测音频中的语音段落
"""

import time
from collections import deque
from typing import Optional, List, Tuple
from dataclasses import dataclass

import numpy as np
import torch

from . import vadconfig as config


@dataclass
class VADResult:
    """VAD 检测结果"""
    has_voice: bool = False          # 当前是否有语音
    voice_start: bool = False        # 是否是语音开始
    voice_end: bool = False          # 是否是语音结束
    probability: float = 0.0         # 语音概率


@dataclass
class SpeechSegment:
    """语音片段"""
    start_ms: float                  # 开始时间 (毫秒)
    end_ms: float                    # 结束时间 (毫秒)


class SileroVAD:
    """
    基于 SileroVAD 的语音活动检测器

    使用方式:
        vad = SileroVAD()
        result = vad.detect(audio_bytes)
        if result.voice_start:
            print("语音开始")
    """

    def __init__(
        self,
        threshold_high: float = None,
        threshold_low: float = None,
        hop_size: int = None,
        frame_window_size: int = None,
        frame_window_threshold: int = None,
        min_silence_duration_ms: int = None,
        model_source: str = None,
    ):
        """
        初始化 SileroVAD

        Args:
            threshold_high: 高阈值
            threshold_low: 低阈值
            hop_size: 帧大小
            frame_window_size: 滑动窗口大小
            frame_window_threshold: 窗口内语音帧阈值
            min_silence_duration_ms: 最小静音时长
            model_source: 模型来源
        """
        # 参数配置
        self.threshold_high = threshold_high or config.THRESHOLD_HIGH
        self.threshold_low = threshold_low or config.THRESHOLD_LOW
        self.hop_size = hop_size or config.HOP_SIZE
        self.frame_window_size = frame_window_size or config.FRAME_WINDOW_SIZE
        self.frame_window_threshold = frame_window_threshold or config.FRAME_WINDOW_THRESHOLD
        self.min_silence_duration_ms = min_silence_duration_ms or config.MIN_SILENCE_DURATION_MS
        model_source = model_source or config.MODEL_SOURCE

        # 状态变量
        self.voice_window = deque(maxlen=self.frame_window_size)
        self.last_is_voice = False
        self.last_activity_time = 0
        self.current_has_voice = False

        # 加载模型
        self.model = None
        self._load_model(model_source)

    def _load_model(self, model_source: str) -> None:
        """加载 SileroVAD 模型"""
        try:
            if model_source.startswith('snakers4/'):
                # 从 GitHub 加载
                self.model, _ = torch.hub.load(
                    repo_or_dir='snakers4/silero-vad',
                    model='silero_vad',
                    force_reload=False,
                )
                print('[VAD] SileroVAD 模型加载成功 (GitHub)')
            else:
                # 从本地加载
                self.model, _ = torch.hub.load(
                    repo_or_dir=model_source,
                    source='local',
                    model='silero_vad',
                    force_reload=False,
                )
                print(f'[VAD] SileroVAD 本地模型加载成功: {model_source}')
        except Exception as e:
            print(f'[VAD Error] 模型加载失败: {e}')
            self.model = None

    def detect(self, audio_data: bytes) -> VADResult:
        """
        检测音频中的语音活动

        Args:
            audio_data: PCM 音频数据 (16bit, 16kHz)

        Returns:
            VADResult: 检测结果
        """
        # 转换为 numpy 数组
        audio_int16 = np.frombuffer(audio_data, dtype=np.int16)

        # 检查数据长度
        if len(audio_int16) < self.hop_size:
            return VADResult(
                has_voice=self.current_has_voice,
                probability=0.0
            )

        voice_start = False
        voice_end = False
        last_probability = 0.0

        # 处理所有完整的帧
        while len(audio_int16) >= self.hop_size:
            current_time = time.time() * 1000

            # 提取一帧
            frame = audio_int16[:self.hop_size]
            audio_int16 = audio_int16[self.hop_size:]

            # 检测单帧
            probability = self._detect_frame(frame)
            last_probability = probability

            # 双阈值判断
            if probability >= self.threshold_high:
                is_voice = True
            elif probability <= self.threshold_low:
                is_voice = False
            else:
                is_voice = self.last_is_voice

            # 更新滑动窗口
            self.voice_window.append(is_voice)
            self.last_is_voice = is_voice

            # 计算窗口内语音帧数
            voice_count = sum(self.voice_window)
            has_voice = (voice_count >= self.frame_window_threshold)

            # 检测状态变化
            if not self.current_has_voice and has_voice:
                voice_start = True
            elif self.current_has_voice and not has_voice:
                silence_duration = current_time - self.last_activity_time
                if silence_duration >= self.min_silence_duration_ms:
                    voice_end = True

            # 更新状态
            if has_voice:
                self.last_activity_time = current_time
            self.current_has_voice = has_voice

        return VADResult(
            has_voice=self.current_has_voice,
            voice_start=voice_start,
            voice_end=voice_end,
            probability=last_probability
        )

    def _detect_frame(self, frame: np.ndarray) -> float:
        """检测单帧语音概率"""
        try:
            if self.model is not None:
                # 归一化到 [-1, 1]
                audio_float = frame.astype(np.float32) / 32768.0
                audio_tensor = torch.from_numpy(audio_float)

                with torch.no_grad():
                    prob = self.model(audio_tensor, config.SAMPLE_RATE).item()
                return prob
            else:
                # 简单能量检测 (备用)
                energy = np.mean(np.abs(frame.astype(np.float32)))
                return min(energy / 1000.0, 1.0)
        except Exception as e:
            print(f'[VAD Error] 检测帧失败: {e}')
            return 0.0

    def detect_file(self, audio_path: str) -> List[SpeechSegment]:
        """
        检测音频文件中的语音段落

        Args:
            audio_path: 音频文件路径 (WAV, 16kHz)

        Returns:
            List[SpeechSegment]: 语音段落列表
        """
        import wave

        with wave.open(audio_path, 'rb') as wav:
            sample_rate = wav.getframerate()
            n_channels = wav.getnchannels()
            n_frames = wav.getnframes()
            audio_data = wav.readframes(n_frames)

        # 如果是双声道，转为单声道
        audio_int16 = np.frombuffer(audio_data, dtype=np.int16)
        if n_channels == 2:
            audio_int16 = audio_int16[::2]

        # 如果采样率不是 16kHz，需要重采样
        if sample_rate != config.SAMPLE_RATE:
            # 简单的线性插值重采样
            ratio = config.SAMPLE_RATE / sample_rate
            new_length = int(len(audio_int16) * ratio)
            indices = np.linspace(0, len(audio_int16) - 1, new_length)
            audio_int16 = np.interp(
                indices,
                np.arange(
                    len(audio_int16)),
                audio_int16).astype(
                np.int16)

        # 重置状态
        self.reset()

        # 逐帧检测
        segments = []
        current_segment_start = None
        ms_per_sample = 1000.0 / config.SAMPLE_RATE

        for i in range(0, len(audio_int16), self.hop_size):
            frame = audio_int16[i:i + self.hop_size]
            if len(frame) < self.hop_size:
                break

            current_ms = i * ms_per_sample
            result = self.detect(frame.tobytes())

            if result.voice_start and current_segment_start is None:
                current_segment_start = current_ms
            elif result.voice_end and current_segment_start is not None:
                segments.append(SpeechSegment(
                    start_ms=current_segment_start,
                    end_ms=current_ms
                ))
                current_segment_start = None

        # 处理最后一个段落
        if current_segment_start is not None:
            end_ms = len(audio_int16) * ms_per_sample
            segments.append(SpeechSegment(
                start_ms=current_segment_start,
                end_ms=end_ms
            ))

        return segments

    def reset(self) -> None:
        """重置 VAD 状态"""
        self.voice_window.clear()
        self.last_is_voice = False
        self.current_has_voice = False
        self.last_activity_time = 0


# ============================================================
# 便捷函数
# ============================================================

def detect_speech_segments(audio_path: str, **kwargs) -> List[SpeechSegment]:
    """
    检测音频文件中的语音段落

    Args:
        audio_path: 音频文件路径
        **kwargs: 传递给 SileroVAD 的参数

    Returns:
        语音段落列表
    """
    vad = SileroVAD(**kwargs)
    return vad.detect_file(audio_path)
