"""Xiaomi MiMo ASR — plugin form.

Subclasses :class:`agent.stt_provider.STTProvider`. Sends audio to MiMo's
chat/completions API with ``input_audio`` format and returns transcribed text.

Xiaomi ASR does **not** use the standard ``/v1/audio/transcriptions`` endpoint.
Audio must be wav or mp3; other formats are auto-converted via ffmpeg.

Auth env var::

    XIAOMI_API_KEY=...               # https://platform.xiaomimimo.com
    XIAOMI_BASE_URL=...              # optional, defaults to https://api.xiaomimimo.com/v1
"""

from __future__ import annotations

import base64
import json
import logging
import os
import ssl
import subprocess
import tempfile
import urllib.error
import urllib.request
from typing import Any, Dict, List, Optional

from agent.stt_provider import STTProvider

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "mimo-v2.5-asr"
DEFAULT_LANGUAGE = "zh"
_DEFAULT_BASE_URL = "https://api.xiaomimimo.com/v1"

_MODELS: List[Dict[str, Any]] = [
    {"id": "mimo-v2.5-asr", "display": "MiMo V2.5 ASR", "languages": ["zh", "en"]},
]


def _get_api_base() -> str:
    return os.environ.get("XIAOMI_BASE_URL", _DEFAULT_BASE_URL).rstrip("/")


class MiMoSTTProvider(STTProvider):
    """Xiaomi MiMo ASR (Speech-to-Text) provider."""

    @property
    def name(self) -> str:
        return "mimo-asr"

    @property
    def display_name(self) -> str:
        return "MiMo ASR (Xiaomi)"

    def is_available(self) -> bool:
        return bool(os.environ.get("XIAOMI_API_KEY"))

    def list_models(self) -> List[Dict[str, Any]]:
        return _MODELS

    def default_model(self) -> Optional[str]:
        return DEFAULT_MODEL

    def default_language(self) -> Optional[str]:
        return DEFAULT_LANGUAGE

    def get_setup_schema(self) -> Dict[str, Any]:
        return {
            "name": "MiMo ASR (Xiaomi)",
            "badge": "paid",
            "tag": "V2.5 — speech-to-text, optimized for Chinese.",
            "env_vars": [
                {
                    "key": "XIAOMI_API_KEY",
                    "prompt": "Xiaomi MiMo API key",
                    "url": "https://platform.xiaomimimo.com",
                },
            ],
        }

    def transcribe(
        self,
        audio_path: str,
        *,
        model: Optional[str] = None,
        language: Optional[str] = None,
        **extra: Any,
    ) -> str:
        api_key = os.environ.get("XIAOMI_API_KEY")
        if not api_key:
            raise RuntimeError(
                "XIAOMI_API_KEY environment variable is not set. "
                "Get your key at https://platform.xiaomimimo.com"
            )

        model_id = model or DEFAULT_MODEL
        model_id = model_id.strip()
        language = language or DEFAULT_LANGUAGE

        audio_path = audio_path.replace("\\", "/")  # normalize Windows paths

        ext = os.path.splitext(audio_path)[1].lower()
        if ext in (".wav",):
            fmt = "wav"
            audio_file = audio_path
        elif ext in (".mp3",):
            fmt = "mp3"
            audio_file = audio_path
        else:
            fmt = "wav"
            try:
                subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
            except (FileNotFoundError, subprocess.CalledProcessError):
                raise RuntimeError(
                    f"ffmpeg is required to convert {ext} to wav/mp3 for Xiaomi ASR"
                )

            fd, audio_file = tempfile.mkstemp(suffix=".wav", prefix="hermes-asr-")
            os.close(fd)
            try:
                subprocess.run(
                    ["ffmpeg", "-y", "-i", audio_path, "-ar", "16000", "-ac", "1",
                     "-f", "wav", audio_file],
                    capture_output=True, check=True,
                )
            except subprocess.CalledProcessError as e:
                os.unlink(audio_file)
                raise RuntimeError(
                    f"ffmpeg conversion failed: {e.stderr.decode(errors='replace')}"
                ) from e

        try:
            with open(audio_file, "rb") as f:
                audio_b64 = base64.b64encode(f.read()).decode()

            payload = json.dumps({
                "model": model_id,
                "messages": [{
                    "role": "user",
                    "content": [
                        {"type": "input_audio", "input_audio": {"data": audio_b64, "format": fmt}},
                    ]
                }]
            }).encode()

            req = urllib.request.Request(
                f"{_get_api_base()}/chat/completions",
                data=payload,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
            )
            ctx = ssl.create_default_context()

            logger.info("MiMo ASR: transcribing %s with model=%s language=%s", audio_path, model_id, language)

            resp = urllib.request.urlopen(req, context=ctx, timeout=300)
            response_data = json.loads(resp.read().decode())

            choices = response_data.get("choices", [])
            if not choices:
                error_info = response_data.get("error", {})
                raise RuntimeError(
                    error_info.get("message", "No transcription result")
                )

            msg = choices[0].get("message", {})
            transcript = msg.get("content", "").strip()

            if not transcript:
                raise RuntimeError("Xiaomi ASR returned empty transcript")

            return transcript

        except urllib.error.HTTPError as e:
            body = e.read().decode(errors="replace")
            raise RuntimeError(f"HTTP {e.code}: {body}") from e
        except urllib.error.URLError as e:
            raise RuntimeError(f"URL error: {e.reason}") from e
        finally:
            if 'audio_file' in locals() and audio_file != audio_path and os.path.exists(audio_file):
                os.unlink(audio_file)