"""
Audio model serving via HuggingFace Transformers.

Exposes:
  POST /v1/audio/transcriptions  → speech-to-text  (ASR models)
  POST /v1/audio/speech          → text-to-speech  (TTS models)
"""

from __future__ import annotations

import io
import os
import struct
import tempfile
import wave
from pathlib import Path

from llmpm import display
from llmpm.core import server as _server

_ASR_TAGS = {"automatic-speech-recognition", "audio-classification"}
_TTS_TAGS = {"text-to-speech", "text-to-audio"}


def load(
    repo_id: str,
    model_dir: Path,
    default_max_tokens: int,
    pipeline_tag: str = "automatic-speech-recognition",
    metadata: dict | None = None,
) -> _server.HandlerContext:
    """Load an audio model and return a HandlerContext."""
    try:
        from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            pipeline as hf_pipeline,
        )
    except ImportError:
        display.error(
            "[bold red]transformers[/] is not installed.\n\n"
            "  Install with:\n\n"
            "    [bold cyan]pip install transformers torch[/]"
        )
        raise SystemExit(1) from None

    display.step(f"Loading {repo_id}…")
    try:
        pipe = hf_pipeline(
            pipeline_tag,
            model=str(model_dir),
            device_map="auto",
            trust_remote_code=True,
        )
    except Exception as exc:  # pylint: disable=broad-except
        display.error(f"Failed to load model: {exc}")
        raise SystemExit(1) from None

    display.ok(f"Model ready  [dim]{repo_id}[/]")

    if pipeline_tag in _ASR_TAGS:
        return _build_asr_ctx(repo_id, pipe, default_max_tokens, metadata or {})
    return _build_tts_ctx(repo_id, pipe, default_max_tokens, metadata or {})


def _build_asr_ctx(
    repo_id: str,
    pipe,  # type: ignore[no-untyped-def]
    default_max_tokens: int,
    metadata: dict | None = None,
) -> _server.HandlerContext:
    def transcribe(audio_bytes: bytes, language: str | None) -> str:
        suffix = ".wav"
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
            f.write(audio_bytes)
            tmp = f.name
        try:
            kwargs = {}
            if language:
                kwargs["generate_kwargs"] = {"language": language}
            result = pipe(tmp, **kwargs)
            return result.get("text", "")
        finally:
            os.unlink(tmp)

    return _server.HandlerContext(
        model_id=repo_id,
        default_max_tokens=default_max_tokens,
        category="automatic-speech-recognition",
        transcribe=transcribe,
        metadata=metadata or {},
    )


def _build_tts_ctx(
    repo_id: str,
    pipe,  # type: ignore[no-untyped-def]
    default_max_tokens: int,
    metadata: dict | None = None,
) -> _server.HandlerContext:
    def synthesize(text: str, voice: str | None, speed: float) -> bytes:
        try:
            import numpy as np  # type: ignore  # pylint: disable=import-outside-toplevel
        except ImportError as exc:
            raise RuntimeError(
                "numpy is required for TTS: pip install numpy"
            ) from exc

        result = pipe(text)
        audio  = result["audio"]
        if hasattr(audio, "numpy"):
            audio = audio.numpy()

        audio_array   = np.array(audio, dtype=np.float32).flatten()
        sampling_rate = result.get("sampling_rate", 22050)
        audio_int16   = (audio_array * 32767).astype(np.int16)

        buf = io.BytesIO()
        with wave.open(buf, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sampling_rate)
            wf.writeframes(
                struct.pack(f"<{len(audio_int16)}h", *audio_int16)
            )
        return buf.getvalue()

    return _server.HandlerContext(
        model_id=repo_id,
        default_max_tokens=default_max_tokens,
        category="text-to-speech",
        synthesize=synthesize,
        metadata=metadata or {},
    )


def serve(
    repo_id: str,
    model_dir: Path,
    host: str,
    port: int,
    default_max_tokens: int,
    pipeline_tag: str = "automatic-speech-recognition",
) -> None:
    """Load an audio model and start the HTTP server."""
    ctx = load(repo_id, model_dir, default_max_tokens, pipeline_tag)
    banner_category = (
        "audio-transcription" if pipeline_tag in _ASR_TAGS else "audio-speech"
    )
    display.serve_banner(host, port, category=banner_category)
    _server.start_multi([ctx], host=host, port=port)
