from __future__ import annotations

import base64
import logging
import struct
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger(__name__)


class AudioFormat(Enum):
    MP3 = "mp3"
    WAV = "wav"
    OGG = "ogg"
    SILK = "silk"
    PCM = "pcm"
    AAC = "aac"
    AMR = "amr"


_AUDIO_MIME_MAP: dict[str, str] = {
    "mp3": "audio/mpeg",
    "wav": "audio/wav",
    "ogg": "audio/ogg",
    "silk": "audio/silk",
    "pcm": "audio/pcm",
    "aac": "audio/aac",
    "amr": "audio/amr",
}

_AUDIO_EXT_MAP: dict[str, str] = {
    "audio/mpeg": ".mp3",
    "audio/wav": ".wav",
    "audio/ogg": ".ogg",
    "audio/silk": ".silk",
    "audio/pcm": ".pcm",
    "audio/aac": ".aac",
    "audio/amr": ".amr",
}

_SUPPORTED_BITRATES: dict[str, list[int]] = {
    "mp3": [8000, 16000, 32000, 64000, 128000],
    "wav": [8000, 16000, 44100],
    "pcm": [8000, 16000, 24000, 44100],
    "aac": [16000, 32000, 64000],
    "ogg": [16000, 32000, 48000],
}


def get_mime_type(fmt: AudioFormat) -> str:
    return _AUDIO_MIME_MAP.get(fmt.value, "audio/mpeg")


def get_extension(fmt: AudioFormat) -> str:
    mime = _AUDIO_MIME_MAP.get(fmt.value, "")
    return _AUDIO_EXT_MAP.get(mime, ".mp3")


def get_supported_bitrates(fmt: AudioFormat) -> list[int]:
    return _SUPPORTED_BITRATES.get(fmt.value, [16000])


def data_uri_to_bytes(data_uri: str) -> tuple[bytes, str]:
    if "," not in data_uri:
        return base64.b64decode(data_uri), "audio/mpeg"

    header, b64_data = data_uri.split(",", 1)
    mime = "audio/mpeg"
    if "data:" in header:
        mime = header.split(":")[1].split(";")[0] if ";" in header.split(":")[1] else header.split(":")[1]

    return base64.b64decode(b64_data), mime


def get_audio_duration_s(audio_data: bytes, fmt: AudioFormat) -> float:
    if fmt == AudioFormat.WAV:
        return _wav_duration(audio_data)
    if fmt == AudioFormat.MP3:
        return _mp3_estimate_duration(audio_data)
    return len(audio_data) / 16000.0


def _wav_duration(data: bytes) -> float:
    try:
        if len(data) < 44:
            return 0.0
        byte_rate = struct.unpack_from("<I", data, 28)[0]
        data_size = struct.unpack_from("<I", data, 40)[0]
        if byte_rate == 0:
            return 0.0
        return data_size / byte_rate
    except (struct.error, IndexError):
        return 0.0


def _mp3_estimate_duration(data: bytes) -> float:
    return len(data) / 16000.0


def calculate_audio_size(
    duration_s: float,
    sample_rate: int = 16000,
    channels: int = 1,
    sample_width: int = 2,
) -> int:
    return int(duration_s * sample_rate * channels * sample_width)


@dataclass
class AudioFormatPolicy:
    transcode_enabled: bool = True
    upload_direct_formats: list[str] = field(default_factory=lambda: ["wav", "mp3"])
    stt_direct_formats: list[str] = field(default_factory=lambda: ["wav", "mp3", "pcm"])
    fallback_format: AudioFormat = field(default=AudioFormat.MP3)
    sample_rate: int = 16000
    channels: int = 1
    bitrate: int = 32000

    def needs_transcode(self, fmt: AudioFormat) -> bool:
        if not self.transcode_enabled:
            return False
        return fmt.value not in self.upload_direct_formats

    def can_stt_direct(self, fmt: AudioFormat) -> bool:
        return fmt.value in self.stt_direct_formats

    @classmethod
    def from_config(cls, config: dict | None) -> AudioFormatPolicy:
        if not config:
            return cls()
        policy_cfg = config.get("audio_format_policy", {})
        return cls(
            transcode_enabled=policy_cfg.get("transcode_enabled", True),
            upload_direct_formats=policy_cfg.get("upload_direct_formats", ["wav", "mp3"]),
            stt_direct_formats=policy_cfg.get("stt_direct_formats", ["wav", "mp3", "pcm"]),
            fallback_format=AudioFormat(policy_cfg.get("fallback_format", "mp3")),
            sample_rate=policy_cfg.get("sample_rate", 16000),
            channels=policy_cfg.get("channels", 1),
            bitrate=policy_cfg.get("bitrate", 32000),
        )


class STTProvider:
    def __init__(
        self,
        provider: str = "",
        api_key: str = "",
        region: str = "",
        model: str = "",
    ):
        self._provider = provider.lower() or "builtin"
        self._api_key = api_key
        self._region = region
        self._model = model

    async def transcribe(self, audio_data: bytes, fmt: AudioFormat | None = None) -> str:
        if self._provider == "azure":
            return await self._transcribe_azure(audio_data)
        elif self._provider == "whisper":
            return await self._transcribe_whisper(audio_data, fmt)
        else:
            return await self._transcribe_builtin(audio_data, fmt)

    async def _transcribe_azure(self, audio_data: bytes) -> str:
        import aiohttp

        try:
            async with aiohttp.ClientSession() as session:
                url = (
                    f"https://{self._region}.stt.speech.microsoft.com/"
                    "speech/recognition/conversation/cognitiveservices/v1"
                    "?language=zh-CN&format=detailed"
                )
                headers = {
                    "Ocp-Apim-Subscription-Key": self._api_key,
                    "Content-Type": "audio/wav",
                }
                async with session.post(url, data=audio_data, headers=headers) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        return data.get("DisplayText", "")
                    logger.warning("STT: Azure returned status %d", resp.status)
        except Exception:
            logger.exception("STT: Azure transcription failed")
        return ""

    async def _transcribe_whisper(self, audio_data: bytes, fmt: AudioFormat | None = None) -> str:
        import json
        import os
        import subprocess
        import tempfile

        ext = ".wav" if fmt is None else get_extension(fmt)
        tmpdir = tempfile.mkdtemp()
        in_path = os.path.join(tmpdir, f"stt_input{ext}")

        try:
            with open(in_path, "wb") as f:
                f.write(audio_data)

            subprocess.run(
                [
                    "whisper",
                    in_path,
                    "--model",
                    self._model or "base",
                    "--output_format",
                    "json",
                    "--output_dir",
                    tmpdir,
                    "--language",
                    "zh",
                ],
                capture_output=True,
                timeout=120,
                check=False,
            )

            json_path = os.path.join(tmpdir, "stt_input.json")
            if os.path.exists(json_path):
                with open(json_path, encoding="utf-8") as f:
                    data = json.load(f)
                return data.get("text", "")
        except Exception:
            logger.exception("STT: Whisper transcription failed")
        finally:
            import shutil

            shutil.rmtree(tmpdir, ignore_errors=True)

        return ""

    async def _transcribe_builtin(self, audio_data: bytes, fmt: AudioFormat | None = None) -> str:
        try:
            import speech_recognition as sr

            ext = ".wav" if fmt is None else get_extension(fmt)
            import os
            import tempfile

            tmpdir = tempfile.mkdtemp()
            in_path = os.path.join(tmpdir, f"stt_input{ext}")

            with open(in_path, "wb") as f:
                f.write(audio_data)

            recognizer = sr.Recognizer()
            with sr.AudioFile(in_path) as source:
                audio = recognizer.record(source)

            import shutil

            shutil.rmtree(tmpdir, ignore_errors=True)

            return recognizer.recognize_google(audio, language="zh-CN")
        except ImportError:
            logger.warning("STT: speech_recognition not installed")
        except Exception:
            logger.exception("STT: builtin transcription failed")
        return ""

    @classmethod
    def from_config(cls, config: dict | None) -> STTProvider:
        if not config:
            return cls()
        stt_cfg = config.get("stt", {})
        return cls(
            provider=stt_cfg.get("provider", "") or config.get("stt_provider", ""),
            api_key=stt_cfg.get("api_key", "") or config.get("stt_api_key", ""),
            region=stt_cfg.get("region", "") or config.get("stt_region", ""),
            model=stt_cfg.get("model", "") or config.get("stt_model", ""),
        )


class TTSProvider:
    def __init__(
        self,
        send_fn: Callable[..., Any] | None = None,
        default_voice: str = "zh-CN-XiaoxiaoNeural",
        default_format: AudioFormat = AudioFormat.MP3,
    ):
        self._send_fn = send_fn
        self._default_voice = default_voice
        self._default_format = default_format
        self._cache: dict[str, bytes] = {}

    async def synthesize(self, text: str, voice: str = "", fmt: AudioFormat | None = None) -> bytes:
        cache_key = f"{text}:{voice}:{fmt.value if fmt else self._default_format.value}"
        cached = self._cache.get(cache_key)
        if cached:
            return cached

        import os

        tts_provider = os.environ.get("QQBOT_TTS_PROVIDER", "builtin").lower()

        if tts_provider == "azure":
            result = await self._synthesize_azure(text, voice or self._default_voice, fmt or self._default_format)
        elif tts_provider == "edge":
            result = await self._synthesize_edge(text, voice or self._default_voice, fmt or self._default_format)
        else:
            result = await self._synthesize_builtin(text)

        self._cache[cache_key] = result
        return result

    async def _synthesize_builtin(self, text: str) -> bytes:
        import os
        import subprocess
        import tempfile

        text_encoded = text.replace('"', '\\"')
        tmpdir = tempfile.mkdtemp()
        out_path = os.path.join(tmpdir, "tts_output.mp3")

        try:
            subprocess.run(
                [
                    "python",
                    "-c",
                    f"import pyttsx3; e=pyttsx3.init(); e.save_to_file('{text_encoded}','{out_path}'); e.runAndWait()",
                ],
                capture_output=True,
                timeout=30,
                check=False,
            )

            if os.path.exists(out_path):
                with open(out_path, "rb") as f:
                    return f.read()
        except Exception:
            logger.exception("TTS: builtin synthesis failed")
        finally:
            import shutil

            shutil.rmtree(tmpdir, ignore_errors=True)

        return self._generate_silence(0.5)

    async def _synthesize_azure(self, text: str, voice: str, fmt: AudioFormat) -> bytes:
        import os

        import aiohttp

        key = os.environ.get("AZURE_TTS_KEY", "")
        region = os.environ.get("AZURE_TTS_REGION", "eastasia")

        if not key:
            logger.warning("TTS: Azure key not configured")
            return self._generate_silence(0.5)

        ssml = (
            f'<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis" xml:lang="zh-CN">'
            f'<voice name="{voice}">{text}</voice></speak>'
        )

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"https://{region}.tts.speech.microsoft.com/cognitiveservices/v1",
                    headers={
                        "Ocp-Apim-Subscription-Key": key,
                        "Content-Type": "application/ssml+xml",
                        "X-Microsoft-OutputFormat": "audio-16khz-32kbitrate-mono-mp3",
                    },
                    data=ssml.encode("utf-8"),
                ) as resp:
                    if resp.status == 200:
                        return await resp.read()
                    logger.warning("TTS: Azure returned status %d", resp.status)
        except Exception:
            logger.exception("TTS: Azure synthesis failed")

        return self._generate_silence(0.5)

    async def _synthesize_edge(self, text: str, voice: str, fmt: AudioFormat) -> bytes:
        try:
            import aiohttp

            ssml = (
                f'<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis" xml:lang="zh-CN">'
                f'<voice name="{voice}">{text}</voice></speak>'
            )

            async with aiohttp.ClientSession() as session:
                async with session.post(
                    "https://speech.platform.bing.com/consumer/speech/synthesize/"
                    "readaloud/edge/v1?TrustedClientToken=6A5AA1D4EAFF4E9FB37E23D68491D6F4",
                    headers={
                        "Content-Type": "application/ssml+xml",
                        "X-Microsoft-OutputFormat": "audio-16khz-32kbitrate-mono-mp3",
                    },
                    data=ssml.encode("utf-8"),
                ) as resp:
                    if resp.status == 200:
                        return await resp.read()
        except Exception:
            logger.exception("TTS: Edge synthesis failed")

        return self._generate_silence(0.5)

    @staticmethod
    def _generate_silence(duration_s: float) -> bytes:
        sample_rate = 16000
        samples = int(duration_s * sample_rate)
        silence = b"\x00" * (samples * 2)
        return silence

    def clear_cache(self) -> None:
        self._cache.clear()