"""
语音识别模块 (STT)

使用腾讯云 ASR 一句话识别。coco 只负责采集/接收音频，识别在云端完成。
"""

from __future__ import annotations

import base64
import tempfile
import threading
import wave
from pathlib import Path

import numpy as np

from ..runtime import tencent_asr_config


TENCENT_ASR_MAX_BYTES = 3 * 1024 * 1024
SUPPORTED_FORMATS = {
    "wav": "wav",
    "pcm": "pcm",
    "ogg": "ogg-opus",
    "opus": "ogg-opus",
    "mp3": "mp3",
    "m4a": "m4a",
    "aac": "aac",
    "amr": "amr",
    "silk": "silk",
    "speex": "speex",
}


class SpeechToText:
    def __init__(
        self,
        region: str = "ap-guangzhou",
        engine_model_type: str = "16k_zh",
        hotword_list: str = "",
        hotword_id: str = "",
        customization_id: str = "",
        filter_dirty: int = 0,
        filter_modal: int = 0,
        filter_punc: int = 0,
        convert_num_mode: int = 1,
        secret_id: str = "",
        secret_key: str = "",
    ):
        self.region = region or "ap-guangzhou"
        self.engine_model_type = engine_model_type or "16k_zh"
        self.hotword_list = hotword_list or ""
        self.hotword_id = hotword_id or ""
        self.customization_id = customization_id or ""
        self.filter_dirty = int(filter_dirty or 0)
        self.filter_modal = int(filter_modal or 0)
        self.filter_punc = int(filter_punc or 0)
        self.convert_num_mode = int(convert_num_mode or 1)
        self.secret_id = secret_id
        self.secret_key = secret_key
        self._client = None
        self.last_audio_stats = {}
        self.last_response = {}

    def transcribe_file(
        self,
        audio_path: str,
        language: str = "zh",
        vad_filter: bool = None,
    ) -> str:
        """识别一个 60 秒以内的短音频文件。"""
        path = Path(audio_path)
        data = path.read_bytes()
        voice_format = self._voice_format_for_path(path)
        self._update_file_stats(path, data)
        return self._recognize_bytes(data, voice_format=voice_format, language=language)

    def transcribe_mic(
        self,
        duration: int = 10,
        sample_rate: int = 16000,
        language: str = "zh",
    ) -> str:
        import sounddevice as sd

        print(f"[录音] {duration} 秒...")
        audio = sd.rec(
            int(duration * sample_rate),
            samplerate=sample_rate,
            channels=1,
            dtype=np.int16,
        )
        sd.wait()
        print("[识别中]")
        return self._save_and_transcribe(audio, sample_rate, language)

    def transcribe_mic_ptt(self, sample_rate: int = 16000, language: str = "zh") -> str:
        import sounddevice as sd

        cmd = input("[按 Enter 开始说话，输入 /q 退出] ").strip()
        if cmd in ("/quit", "/q"):
            return "/QUIT"

        chunks = []
        stop_event = threading.Event()

        def callback(indata, frames, time_info, status):
            chunks.append(indata.copy())

        def wait_for_stop():
            input()
            stop_event.set()

        reader = threading.Thread(target=wait_for_stop, daemon=True)
        stream = sd.InputStream(
            samplerate=sample_rate,
            channels=1,
            dtype=np.int16,
            callback=callback,
        )
        stream.start()
        print("[聆听中...说完按 Enter 结束]")
        reader.start()

        stop_event.wait()
        stream.stop()
        stream.close()
        reader.join(timeout=1)

        if not chunks:
            return ""

        audio = np.concatenate(chunks, axis=0)
        dur = len(audio) / sample_rate
        print(f"[已录制 {dur:.1f} 秒，识别中...]")
        return self._save_and_transcribe(audio, sample_rate, language)

    def save_and_transcribe_array(
        self,
        audio,
        sample_rate: int = 16000,
        language: str = "zh",
        vad_filter: bool = None,
    ) -> str:
        """识别内存中的 PCM int16 音频。"""
        return self._save_and_transcribe(audio, sample_rate, language)

    def _save_and_transcribe(self, audio, sample_rate: int, language: str) -> str:
        audio = self._prepare_audio(audio, sample_rate=sample_rate)
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            wav_path = Path(f.name)
        try:
            with wave.open(str(wav_path), "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)
                wf.setframerate(sample_rate)
                wf.writeframes(audio.tobytes())
            return self.transcribe_file(str(wav_path), language=language)
        finally:
            wav_path.unlink(missing_ok=True)

    def _recognize_bytes(self, data: bytes, voice_format: str, language: str = "zh") -> str:
        if not data:
            return ""
        if len(data) > TENCENT_ASR_MAX_BYTES:
            raise ValueError("腾讯云一句话识别音频不能超过 3MB")

        client, models = self._tencent_client()
        req = models.SentenceRecognitionRequest()
        req.EngSerViceType = self.engine_model_type
        req.SourceType = 1
        req.VoiceFormat = voice_format
        req.ProjectId = 0
        req.SubServiceType = 2
        req.Data = base64.b64encode(data).decode("ascii")
        req.DataLen = len(data)
        req.FilterDirty = self.filter_dirty
        req.FilterModal = self.filter_modal
        req.FilterPunc = self.filter_punc
        req.ConvertNumMode = self.convert_num_mode
        if self.hotword_list:
            req.HotwordList = self.hotword_list
        if self.hotword_id:
            req.HotwordId = self.hotword_id
        if self.customization_id:
            req.CustomizationId = self.customization_id

        response = client.SentenceRecognition(req)
        self.last_response = {
            "request_id": getattr(response, "RequestId", ""),
            "audio_duration": getattr(response, "AudioDuration", None),
            "engine": self.engine_model_type,
            "format": voice_format,
        }
        return (getattr(response, "Result", "") or "").strip()

    def _tencent_client(self):
        if self._client is not None:
            from tencentcloud.asr.v20190614 import models
            return self._client, models

        try:
            from tencentcloud.common import credential
            from tencentcloud.common.profile.client_profile import ClientProfile
            from tencentcloud.common.profile.http_profile import HttpProfile
            from tencentcloud.asr.v20190614 import asr_client, models
        except ImportError as exc:
            raise RuntimeError(
                "缺少腾讯云 ASR SDK，请运行: pip install tencentcloud-sdk-python-asr"
            ) from exc

        config = tencent_asr_config()
        secret_id = self.secret_id or config.get("secret_id", "")
        secret_key = self.secret_key or config.get("secret_key", "")
        region = self.region or config.get("region", "") or "ap-guangzhou"
        engine = config.get("engine_model_type", "")
        if engine and self.engine_model_type == "16k_zh":
            self.engine_model_type = engine

        if not secret_id or not secret_key:
            raise RuntimeError(
                "未配置腾讯云 ASR 密钥。请设置环境变量 TENCENTCLOUD_SECRET_ID/"
                "TENCENTCLOUD_SECRET_KEY，或写入 .coco/config.yaml 的 tencent_asr。"
            )

        cred = credential.Credential(secret_id, secret_key)
        http_profile = HttpProfile()
        http_profile.endpoint = "asr.tencentcloudapi.com"
        client_profile = ClientProfile()
        client_profile.httpProfile = http_profile
        self._client = asr_client.AsrClient(cred, region, client_profile)
        return self._client, models

    def _voice_format_for_path(self, path: Path) -> str:
        suffix = path.suffix.lower().lstrip(".")
        if suffix == "webm":
            raise ValueError("腾讯云 ASR 不支持 webm，请上传 wav/mp3/m4a/ogg-opus 音频")
        return SUPPORTED_FORMATS.get(suffix, "wav")

    def _update_file_stats(self, path: Path, data: bytes) -> None:
        stats = {
            "bytes": len(data),
            "format": path.suffix.lower().lstrip(".") or "wav",
        }
        if path.suffix.lower() == ".wav":
            try:
                with wave.open(str(path), "rb") as wf:
                    frames = wf.getnframes()
                    rate = wf.getframerate()
                    stats["duration"] = round(frames / max(rate, 1), 3)
                    stats["sample_rate"] = rate
                    stats["channels"] = wf.getnchannels()
                    if wf.getsampwidth() == 2 and frames:
                        samples = np.frombuffer(wf.readframes(frames), dtype=np.int16)
                        if stats["channels"] > 1:
                            samples = samples.reshape(-1, stats["channels"]).mean(axis=1)
                        normalized = samples.astype(np.float32) / 32768.0
                        stats["rms"] = round(float(np.sqrt(np.mean(normalized ** 2))), 5)
                        stats["peak"] = round(float(np.max(np.abs(normalized))), 5)
            except wave.Error:
                pass
        self.last_audio_stats = stats

    def _prepare_audio(self, audio, sample_rate: int = 16000) -> np.ndarray:
        """转单声道、去直流偏移、裁掉长静音，并在音量偏低时自动增益。"""
        arr = np.asarray(audio)
        if arr.ndim > 1:
            arr = arr.mean(axis=1)

        arr = arr.astype(np.float32)
        if arr.size == 0:
            return arr.astype(np.int16)

        if float(np.max(np.abs(arr))) <= 1.5:
            arr *= 32767.0
        trimmed = self._trim_silence(arr, sample_rate)
        if trimmed.size:
            arr = trimmed
        arr -= float(np.mean(arr))

        peak = float(np.max(np.abs(arr)))
        rms = float(np.sqrt(np.mean((arr / 32768.0) ** 2))) if arr.size else 0.0
        gain = 1.0
        if peak > 0 and peak < 0.35 * 32767:
            gain = min((0.85 * 32767) / peak, 8.0)
            arr *= gain

        self.last_audio_stats = {
            "duration": round(float(arr.size) / max(sample_rate, 1), 3),
            "rms": round(rms, 5),
            "peak": round(peak / 32768.0, 5),
            "gain": round(gain, 3),
            "sample_rate": sample_rate,
        }
        return np.clip(arr, -32768, 32767).astype(np.int16)

    def _trim_silence(self, audio: np.ndarray, sample_rate: int) -> np.ndarray:
        if sample_rate <= 0 or audio.size < int(sample_rate * 0.4):
            return audio

        frame = max(1, int(sample_rate * 0.03))
        hop = frame
        rms_values = []
        for start in range(0, max(audio.size - frame + 1, 0), hop):
            chunk = audio[start:start + frame] / 32768.0
            rms_values.append(float(np.sqrt(np.mean(chunk ** 2))))

        if not rms_values:
            return audio
        threshold = max(0.004, min(0.03, max(rms_values) * 0.18))
        voiced = [i for i, value in enumerate(rms_values) if value >= threshold]
        if not voiced:
            return audio

        pad = int(0.18 * sample_rate)
        start = max(0, voiced[0] * hop - pad)
        end = min(audio.size, (voiced[-1] * hop + frame) + pad)
        return audio[start:end]
