"""
core/stt.py - 语音转文字 (STT) 引擎
===================================
[v1.22.0] 从 api_server.py 提取的通用 STT 模块。
支持 4 种引擎按优先级自动选择:
  1. SenseVoice (阿里达摩院, 中文最佳)
  2. vosk (离线本地)
  3. LLM API Whisper 兼容端点 (无额外依赖)
  4. SpeechRecognition (Google, 需外网)
"""
from __future__ import annotations

import io
import os
import re
from typing import Any, Dict, Optional

from core.logger import get_logger

logger = get_logger("myagent.stt")

# 全局模型缓存
_sensevoice_model = None
_vosk_model = None
_llm_config = None  # (api_key, base_url) 元组，用于 Whisper API


def set_llm_config(api_key: str, base_url: str):
    """设置 LLM API 配置（用于 Whisper 兼容端点）"""
    global _llm_config
    _llm_config = (api_key, base_url)


def _convert_to_wav(audio_data: bytes, audio_format: Optional[str] = None) -> bytes:
    """将音频数据转换为 16kHz 单声道 WAV"""
    try:
        from pydub import AudioSegment
        audio_buf = io.BytesIO(audio_data)
        seg = AudioSegment.from_file(audio_buf, format=audio_format or "webm")
        # [v1.23.2] 检查音频时长，过短直接返回原始数据
        if seg.duration_seconds < 0.1:
            logger.debug(f"音频过短 ({seg.duration_seconds:.2f}s)，跳过转换")
            return audio_data
        seg = seg.set_channels(1).set_frame_rate(16000).set_sample_width(2)
        wav_buf = io.BytesIO()
        seg.export(wav_buf, format="wav")
        wav_buf.seek(0)
        return wav_buf.read()
    except Exception as e:
        import shutil
        if not shutil.which("ffmpeg"):
            logger.warning(f"pydub 转换失败且缺少 ffmpeg: {e}")
        else:
            logger.warning(f"pydub 音频转换失败: {e}")
        return audio_data


async def _stt_sensevoice(audio_data: bytes, audio_format: Optional[str] = None) -> Optional[Dict]:
    """SenseVoice 引擎（阿里达摩院，中文识别最佳）"""
    global _sensevoice_model
    try:
        os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
        from funasr import AutoModel

        if _sensevoice_model is None:
            model_dir = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                'models', 'sensevoice',
            )
            # [v1.23.49] 首次加载：将 ModelScope 缓存复制到本地 model_dir，后续不再联网
            if not os.path.isdir(model_dir) or not os.listdir(model_dir):
                import shutil
                ms_cache = os.path.expanduser("~/.cache/modelscope/hub/models/iic/SenseVoiceSmall")
                if os.path.isdir(ms_cache) and os.listdir(ms_cache):
                    os.makedirs(model_dir, exist_ok=True)
                    logger.info(f"复制 SenseVoice 模型: {ms_cache} -> {model_dir}")
                    # 复制模型文件（排除缓存元数据）
                    for item in os.listdir(ms_cache):
                        if item.endswith('.tmp') or item in ('.lock', '__pycache__'):
                            continue
                        src = os.path.join(ms_cache, item)
                        dst = os.path.join(model_dir, item)
                        if os.path.isdir(src):
                            if os.path.exists(dst):
                                shutil.rmtree(dst)
                            shutil.copytree(src, dst)
                        else:
                            shutil.copy2(src, dst)
                else:
                    # ModelScope 缓存也没有，首次下载到本地
                    os.makedirs(model_dir, exist_ok=True)
                    logger.info("SenseVoice 模型首次下载到本地...")
            _sensevoice_model = AutoModel(
                model="iic/SenseVoiceSmall",
                model_dir=model_dir,
                device="cpu",
                disable_pbar=True,
                disable_update=True,
            )
            logger.info("SenseVoice 模型已加载 (iic/SenseVoiceSmall, CPU, 本地缓存)")

        # [v1.23.2] 增强: pydub 转换失败记录警告、WAV 头验证、音频长度检查
        wav_data = _convert_to_wav(audio_data, audio_format)
        wav_path = f"/tmp/myagent_stt_{id(audio_data) % 100000}.wav"
        try:
            # 验证 WAV 文件头 (RIFF....WAVE)
            if len(wav_data) < 44 or wav_data[:4] != b'RIFF' or wav_data[8:12] != b'WAVE':
                logger.warning(f"SenseVoice 跳过: 无效 WAV 数据 (size={len(wav_data)}, header={wav_data[:12].hex()})")
                return None

            with open(wav_path, 'wb') as f:
                f.write(wav_data)

            res = _sensevoice_model.generate(
                input=wav_path,
                cache={},
                language="auto",
                use_itn=True,
                batch_size_s=300,
            )
            if res and len(res) > 0 and len(res[0]) > 0:
                text = res[0][0]["text"] if isinstance(res[0][0], dict) else str(res[0][0])
                # 清理 SenseVoice 特殊 token
                text = re.sub(r'<\|[^|]+\|>', '', text).strip()
                if text:
                    return {"success": True, "output": text, "engine": "sensevoice"}
        finally:
            try:
                os.remove(wav_path)
            except Exception:
                pass
    except ImportError:
        logger.debug("SenseVoice (funasr) 未安装，跳过。安装: pip install funasr torch torchaudio")
    except Exception as e:
        err_str = str(e)
        if "ffmpeg" in err_str.lower() or "No such file" in err_str:
            logger.warning(f"SenseVoice 失败 (缺少 ffmpeg): {e}")
        else:
            logger.warning(f"SenseVoice 失败: {e}")
    return None


async def _stt_vosk(audio_data: bytes) -> Optional[Dict]:
    """vosk 引擎（离线本地）"""
    global _vosk_model
    try:
        import vosk
        import json as _json

        if _vosk_model is None:
            import zipfile
            model_dir = os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                'models', 'vosk', 'vosk-model-cn',
            )
            if not os.path.exists(model_dir):
                logger.info("正在下载 vosk 中文模型...")
                import urllib.request
                url = "https://alphacephei.com/vosk/models/vosk-model-small-cn-0.22.zip"
                zip_path = model_dir + ".zip"
                os.makedirs(os.path.dirname(model_dir), exist_ok=True)
                try:
                    urllib.request.urlretrieve(url, zip_path)
                    with zipfile.ZipFile(zip_path, 'r') as zf:
                        zf.extractall(os.path.dirname(model_dir))
                    os.remove(zip_path)
                except Exception as de:
                    logger.warning(f"vosk 模型下载失败: {de}")
            if os.path.exists(model_dir):
                _vosk_model = vosk.Model(model_dir)
            else:
                return None

        if not _vosk_model:
            return None

        rec = vosk.KaldiRecognizer(_vosk_model, 16000)
        rec.AcceptWaveform(audio_data)
        result = _json.loads(rec.Result())
        text = result.get("text", "").strip()
        if text:
            return {"success": True, "output": text, "engine": "vosk"}
    except ImportError:
        logger.debug("vosk 未安装，跳过")
    except Exception as e:
        logger.warning(f"vosk 转录失败: {e}")
    return None


async def _stt_whisper_api(audio_data: bytes, audio_format: Optional[str] = None) -> Optional[Dict]:
    """LLM API Whisper 兼容端点（无额外依赖）"""
    global _llm_config
    if not _llm_config:
        return None
    api_key, base_url = _llm_config
    if not api_key or not base_url:
        return None

    try:
        import aiohttp
        import mimetypes

        # 构造 Whisper API URL
        base = base_url.rstrip("/")
        if base.endswith("/v1"):
            whisper_url = base + "/audio/transcriptions"
        else:
            whisper_url = base.rstrip("/v1") + "/v1/audio/transcriptions"

        fmt = audio_format or "wav"
        mime = mimetypes.guess_type(f"audio.{fmt}")[0] or "audio/wav"

        data = aiohttp.FormData()
        data.add_field('file', audio_data, filename=f'audio.{fmt}', content_type=mime)
        data.add_field('model', 'whisper-1')
        data.add_field('language', 'zh')

        headers = {"Authorization": f"Bearer {api_key}"}
        async with aiohttp.ClientSession() as session:
            async with session.post(
                whisper_url, data=data, headers=headers,
                timeout=aiohttp.ClientTimeout(total=30),
            ) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    text = result.get("text", "").strip()
                    if text:
                        logger.info(f"LLM API (Whisper) 转录成功")
                        return {"success": True, "output": text, "engine": "llm_api"}
                else:
                    err_text = await resp.text()
                    logger.debug(f"Whisper 端点不可用 ({resp.status}): {err_text[:200]}")
    except Exception as e:
        logger.debug(f"LLM API Whisper 转录失败: {e}")
    return None


async def _stt_speech_recognition(audio_data: bytes, audio_format: Optional[str] = None) -> Optional[Dict]:
    """SpeechRecognition (Google Web Speech API，需外网)"""
    try:
        import speech_recognition as sr

        wav_data = _convert_to_wav(audio_data, audio_format)
        wav_buf = io.BytesIO(wav_data)
        wav_buf.seek(0)

        recognizer = sr.Recognizer()
        with sr.AudioFile(wav_buf) as source:
            audio = recognizer.record(source)
        text = recognizer.recognize_google(audio, language="zh-CN")
        if text:
            logger.info("SpeechRecognition (Google API) 转录成功")
            return {"success": True, "output": text, "engine": "speech_recognition"}
    except ImportError:
        logger.debug("SpeechRecognition 未安装，跳过")
    except sr.UnknownValueError:
        logger.debug("SpeechRecognition 无法识别音频内容")
    except sr.RequestError as e:
        logger.warning(f"SpeechRecognition API 请求失败: {e}")
    except Exception as e:
        logger.warning(f"SpeechRecognition 转录失败: {e}")
    return None


async def transcribe(
    audio_path: str = "",
    audio_data: Optional[bytes] = None,
    audio_format: Optional[str] = None,
    language: str = "zh",
) -> Dict[str, Any]:
    """
    语音转文字 — 使用最佳可用引擎。

    Args:
        audio_path: 音频文件路径（与 audio_data 二选一）
        audio_data: 音频二进制数据（与 audio_path 二选一）
        audio_format: 音频格式 (wav/webm/ogg/mp3 等)，不传则从文件名推断
        language: 语言代码 (zh/en 等)，默认 zh

    Returns:
        {"success": bool, "output": str, "engine": str, "error": str}
    """
    # 读取文件
    if audio_path and not audio_data:
        p = _P(audio_path).expanduser().resolve()
        if not p.exists():
            return {"success": False, "error": f"文件不存在: {audio_path}"}
        audio_data = p.read_bytes()
        # 从文件名推断格式
        ext = p.suffix.lower().lstrip(".")
        if not audio_format and ext:
            audio_format = ext

    if not audio_data:
        return {"success": False, "error": "未收到音频数据"}

    # 4 引擎依次尝试
    engines = [
        ("SenseVoice", lambda: _stt_sensevoice(audio_data, audio_format)),
        ("vosk", lambda: _stt_vosk(audio_data)),
        ("Whisper API", lambda: _stt_whisper_api(audio_data, audio_format)),
        ("SpeechRecognition", lambda: _stt_speech_recognition(audio_data, audio_format)),
    ]

    for name, fn in engines:
        result = await fn()
        if result and result.get("success"):
            return result

    return {
        "success": False,
        "error": (
            "未检测到可用的 STT 引擎。请尝试以下方案:\n"
            "  1. pip install funasr torch torchaudio  (SenseVoice，中文最佳，推荐)\n"
            "  2. 配置支持 Whisper 的 LLM API（自动使用，无需安装）\n"
            "  3. pip install vosk             (离线本地，需下载模型)\n"
            "  4. pip install SpeechRecognition (需外网)"
        ),
    }