"""独立实时 Demo 使用的 OpenAI 兼容 VLLM 服务适配器。"""

from __future__ import annotations

import io
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from aiohttp import ClientSession, ClientTimeout, FormData


@dataclass(frozen=True)
class ModelServiceConfig:
    """一个独立 VLLM 端点所需的连接配置。"""

    base_url: str = "http://127.0.0.1:9950/v1"
    model: str = "Qwen/Qwen3-ASR-0.6B"
    api_key: str = "EMPTY"
    timeout_seconds: float = 45.0


def pcm16_to_wav(pcm_bytes: bytes, sample_rate: int = 16000) -> bytes:
    """将浏览器发送的 PCM16 单声道数据封装为 VLLM 可识别的 WAV 请求。"""
    output = io.BytesIO()
    with wave.open(output, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(pcm_bytes)
    return output.getvalue()


def wav_to_pcm16(audio_bytes: bytes) -> bytes:
    """从 WAV 缓冲区提取 PCM 帧，并兼容尚未完整的中间音频数据。"""
    try:
        with wave.open(io.BytesIO(audio_bytes), "rb") as wav_file:
            return wav_file.readframes(wav_file.getnframes())
    except (EOFError, wave.Error):
        if audio_bytes[:4] == b"RIFF" and audio_bytes[8:12] == b"WAVE" and len(audio_bytes) > 44:
            return audio_bytes[44:]
        return audio_bytes


def prepare_audio_request(
    audio_bytes: bytes,
    source: str,
    file_name: str,
    partial: bool,
) -> tuple[bytes, str, str] | None:
    """将麦克风、PCM 或 WAV 数据转换为 WAV；压缩格式的中间片段延迟到最终帧处理。"""
    suffix = Path(file_name).suffix.lower()
    if source == "mic" or suffix in {".pcm", ".wav"}:
        pcm_bytes = wav_to_pcm16(audio_bytes) if suffix == ".wav" else audio_bytes
        return pcm16_to_wav(pcm_bytes), "audio.wav", "audio/wav"
    if partial:
        # MP3/M4A/OGG 的不断增长前缀通常不是完整容器，不能安全解码，因此只在
        # 最终阶段提交压缩文件，避免中间请求产生随机解码错误。
        return None
    content_type = {
        ".mp3": "audio/mpeg",
        ".m4a": "audio/mp4",
        ".ogg": "audio/ogg",
        ".opus": "audio/ogg",
    }.get(suffix, "application/octet-stream")
    return audio_bytes, Path(file_name).name or "audio.bin", content_type


class VLLMTranscriptionService:
    """只调用独立项目提供的 VLLM HTTP 接口，不导入原项目应用代码。"""

    native_partial_supported = False

    def __init__(self, config: ModelServiceConfig) -> None:
        self.config = config
        self._session: ClientSession | None = None

    async def start(self) -> None:
        """创建可复用的 HTTP 会话，供所有中间和最终转写请求共享。"""
        self._session = ClientSession(timeout=ClientTimeout(total=self.config.timeout_seconds))

    async def close(self) -> None:
        """本地 Demo 退出时释放可复用的 HTTP 会话和底层连接。"""
        if self._session is not None:
            await self._session.close()
            self._session = None

    async def transcribe(
        self,
        audio_bytes: bytes,
        source: str,
        file_name: str,
        partial: bool,
    ) -> str | None:
        """提交一次音频快照并返回文本；返回 None 表示当前格式不支持中间转写。"""
        prepared = prepare_audio_request(audio_bytes, source, file_name, partial)
        if prepared is None:
            return None
        payload, upload_name, content_type = prepared
        if self._session is None:
            raise RuntimeError("model service is not started")

        form = FormData()
        form.add_field("file", payload, filename=upload_name, content_type=content_type)
        form.add_field("model", self.config.model)
        form.add_field("response_format", "json")
        headers = {"Authorization": f"Bearer {self.config.api_key}"}
        endpoint = self.config.base_url.rstrip("/") + "/audio/transcriptions"
        async with self._session.post(endpoint, data=form, headers=headers) as response:
            body = await response.text()
            if response.status >= 400:
                raise RuntimeError(f"VLLM transcription failed ({response.status}): {body[:500]}")
            try:
                decoded: Any = await response.json(content_type=None)
            except ValueError:
                return body.strip()
            if isinstance(decoded, dict):
                return str(decoded.get("text") or decoded.get("transcript") or "").strip()
            return str(decoded).strip()