"""MiMo ASR adapter — ``POST /v1/chat/completions`` with ``input_audio``.

MiMo 的 ASR 不走独立音频端点，而是通过聊天补全接口扩展实现：

    POST {base_url}/chat/completions
    Content-Type: application/json
    Authorization: Bearer {api_key}

    {
        "model": "mimo-v2.5-asr",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "input_audio",
                        "input_audio": {
                            "data": "data:audio/wav;base64,<base64_encoded_audio>"
                        }
                    }
                ]
            }
        ],
        "asr_options": {
            "language": "auto"
        }
    }

响应：:

    {
      "choices": [{
        "message": {
          "content": "转写后的文本"
        }
      }]
    }
"""

from __future__ import annotations

import base64
import hashlib
import time

import httpx

from app.ai_provider.audio.base import (
    ASRProtocol,
    ASRRequest,
    ASRResponse,
    AudioAdapterBase,
    AudioServiceError,
)
from app.core.logging import logger

_MIME_MAP = {
    "wav": "audio/wav",
    "mp3": "audio/mpeg",
    "m4a": "audio/mp4",
    "flac": "audio/flac",
}

# ASR 缓存：默认最多 128 条，TTL 1 小时
_ASR_CACHE_MAX_SIZE = 128
_ASR_CACHE_TTL_SECONDS = 3600


class _ASRCacheEntry:
    __slots__ = ("text", "stored_at")

    def __init__(self, text: str) -> None:
        self.text = text
        self.stored_at = time.monotonic()


_asr_cache: dict[str, _ASRCacheEntry] = {}


def _asr_cache_key(model: str, audio_data: bytes, language: str, fmt: str) -> str:
    """生成 ASR 缓存键：SHA256(model + audio_bytes + language + format)。"""
    raw = f"{model}\0{language}\0{fmt}\0"
    raw_bytes = raw.encode() + audio_data
    return hashlib.sha256(raw_bytes).hexdigest()


def _asr_cache_get(key: str) -> str | None:
    entry = _asr_cache.get(key)
    if entry is None:
        return None
    if time.monotonic() - entry.stored_at > _ASR_CACHE_TTL_SECONDS:
        del _asr_cache[key]
        return None
    return entry.text


def _asr_cache_set(key: str, text: str) -> None:
    if len(_asr_cache) >= _ASR_CACHE_MAX_SIZE:
        sorted_keys = sorted(_asr_cache.keys(), key=lambda k: _asr_cache[k].stored_at)
        for old_key in sorted_keys[: _ASR_CACHE_MAX_SIZE // 2]:
            del _asr_cache[old_key]
    _asr_cache[key] = _ASRCacheEntry(text)


class HunjingASROpenAI(AudioAdapterBase, ASRProtocol):
    """MiMo ASR 适配器，通过 ``POST /chat/completions`` 携带 ``input_audio`` 实现。

    将音频字节包装为 chat 消息中的 ``input_audio`` 内容类型的 data URL，
    然后从响应中提取转写文本。
    """

    _adapter_type = "asr"
    _default_timeout = 300

    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        model: str = "mimo-v2.5-asr",
        request_timeout_seconds: int = 300,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._model = model
        self._timeout = request_timeout_seconds
        self._client: httpx.AsyncClient | None = None

    # ── 客户端 ───────────────────────────────────────────────────────────

    @property
    def _http(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
        return self._client

    async def close(self) -> None:
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()

    # ── ASRProtocol ──────────────────────────────────────────────────────

    async def transcribe(self, request: ASRRequest) -> ASRResponse:
        """通过 ``POST /chat/completions`` 转写音频。

        将音频编码为 data URL 并以 ``input_audio`` 内容块的形式发送到 chat 消息中。
        """
        # ── 缓存查找 ──────────────────────────────────────────────────
        cache_key = _asr_cache_key(
            self._model, request.audio_data, request.language or "auto", request.format
        )
        cached = _asr_cache_get(cache_key)
        if cached is not None:
            logger.info("audio_asr.cache_hit", model=self._model, text_length=len(cached))
            return ASRResponse(text=cached)

        # ── 构造请求 ──────────────────────────────────────────────────
        mime = _MIME_MAP.get(request.format, "audio/wav")
        b64_audio = base64.b64encode(request.audio_data).decode()
        data_url = f"data:{mime};base64,{b64_audio}"

        payload: dict = {
            "model": self._model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_audio",
                            "input_audio": {
                                "data": data_url,
                            },
                        },
                    ],
                },
            ],
            "asr_options": {
                "language": request.language or "auto",
            },
        }

        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

        url = f"{self._base_url}/chat/completions"

        logger.info(
            "audio_asr.request",
            model=self._model,
            format=request.format,
            language=request.language,
            url=url,
            size_bytes=len(request.audio_data),
        )

        # ── API 调用 ────────────────────────────────────────────────
        try:
            response = await self._http.post(url, json=payload, headers=headers)
        except httpx.TimeoutException as exc:
            raise AudioServiceError(
                status_code=504,
                code="audio_asr_timeout",
                message="ASR request timed out calling chat completions.",
            ) from exc
        except httpx.RequestError as exc:
            raise AudioServiceError(
                status_code=503,
                code="audio_asr_connection_error",
                message=f"ASR connection error: {exc}",
            ) from exc

        if response.status_code != 200:
            body = response.text[:500]
            raise AudioServiceError(
                status_code=response.status_code,
                code="audio_asr_api_error",
                message=f"ASR API returned {response.status_code}: {body}",
            )

        # 解析响应，提取转写文本
        text = self._extract_text(response.json())
        if not text:
            raise AudioServiceError(
                status_code=500,
                code="audio_asr_empty_response",
                message="ASR API returned empty transcription.",
            )

        # ── 写入缓存 ────────────────────────────────────────────────
        _asr_cache_set(cache_key, text)

        logger.info(
            "audio_asr.success",
            model=self._model,
            text_length=len(text),
            cache_size=len(_asr_cache),
        )
        return ASRResponse(text=text)

    # ── 辅助方法 ──────────────────────────────────────────────────────────

    @staticmethod
    def _extract_text(body: dict) -> str:
        """从 chat completion 响应中提取转写文本。"""
        try:
            choices = body.get("choices", [])
            if not choices:
                raise AudioServiceError(
                    status_code=500,
                    code="audio_asr_no_choices",
                    message="ASR response has no choices.",
                )
            return (choices[0].get("message") or {}).get("content") or ""
        except AudioServiceError:
            raise
        except Exception as exc:
            raise AudioServiceError(
                status_code=500,
                code="audio_asr_response_parse_error",
                message=f"Failed to parse ASR response: {exc}",
            ) from exc