from __future__ import annotations

import asyncio
import io
import json
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

import httpx
import numpy as np

from .models import ASRStart
from .settings import ASRSettings, ConfigurationError


class ProviderError(RuntimeError):
    pass


@dataclass(frozen=True, slots=True)
class Transcript:
    text: str
    language: str | None = None
    confidence: float | None = None


class ASRProvider(Protocol):
    name: str
    supports_partials: bool

    async def partial(self, pcm: bytes, start: ASRStart) -> Transcript | None: ...
    async def final(self, pcm: bytes, start: ASRStart) -> Transcript: ...
    async def close(self) -> None: ...


def pcm_s16le_to_wav(pcm: bytes, start: ASRStart) -> bytes:
    frame_width = 2 * start.channels
    if len(pcm) % frame_width:
        raise ProviderError("PCM payload is not aligned to complete sample frames")
    output = io.BytesIO()
    with wave.open(output, "wb") as wav:
        wav.setnchannels(start.channels)
        wav.setsampwidth(2)
        wav.setframerate(start.sample_rate_hz)
        wav.writeframes(pcm)
    return output.getvalue()


class MockASRProvider:
    """Deterministic fake. It never opens a socket or loads a model."""

    name = "mock"
    supports_partials = True

    def __init__(
        self,
        *,
        final_text: str = "Hallo von OMNI",
        partial_text: str = "Hallo",
        partial_after_bytes: int = 640,
        language: str = "de",
        confidence: float = 0.99,
    ) -> None:
        self.final_text = final_text
        self.partial_text = partial_text
        self.partial_after_bytes = partial_after_bytes
        self.language = language
        self.confidence = confidence
        self.final_calls = 0
        self.partial_calls = 0

    async def partial(self, pcm: bytes, start: ASRStart) -> Transcript | None:
        self.partial_calls += 1
        if len(pcm) < self.partial_after_bytes:
            return None
        return Transcript(self.partial_text, start.language or self.language, self.confidence)

    async def final(self, pcm: bytes, start: ASRStart) -> Transcript:
        self.final_calls += 1
        if not pcm:
            raise ProviderError("No audio received")
        return Transcript(self.final_text, start.language or self.language, self.confidence)

    async def close(self) -> None:
        return None


class _LocalMultipartProvider:
    supports_partials = False

    def __init__(
        self,
        *,
        endpoint: str,
        model: str,
        api_key: str | None,
        connect_timeout_s: float,
        write_timeout_s: float,
        transcribe_timeout_s: float,
        max_response_bytes: int,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self.endpoint = endpoint
        self.model = model
        self.max_response_bytes = max_response_bytes
        self._owns_client = client is None
        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
        self._client = client or httpx.AsyncClient(
            headers=headers,
            timeout=httpx.Timeout(
                connect=connect_timeout_s,
                read=transcribe_timeout_s,
                write=write_timeout_s,
                pool=connect_timeout_s,
            ),
            follow_redirects=False,
        )

    async def partial(self, pcm: bytes, start: ASRStart) -> None:
        # Re-transcribing growing buffers would compete with MiniCPM and increase latency.
        return None

    def form_fields(self, start: ASRStart) -> dict[str, str]:
        fields = {"model": self.model, "response_format": "json"}
        if start.language:
            fields["language"] = start.language
        return fields

    async def final(self, pcm: bytes, start: ASRStart) -> Transcript:
        wav = pcm_s16le_to_wav(pcm, start)
        try:
            async with self._client.stream(
                "POST",
                self.endpoint,
                data=self.form_fields(start),
                files={"file": ("audio.wav", wav, "audio/wav")},
            ) as response:
                if response.status_code < 200 or response.status_code >= 300:
                    raise ProviderError(f"ASR endpoint returned HTTP {response.status_code}")
                body = bytearray()
                async for chunk in response.aiter_bytes():
                    if len(body) + len(chunk) > self.max_response_bytes:
                        raise ProviderError("ASR response exceeded its configured size limit")
                    body.extend(chunk)
        except ProviderError:
            raise
        except (httpx.TimeoutException, httpx.NetworkError) as exc:
            raise ProviderError("ASR endpoint was unavailable or timed out") from exc
        return self._parse_response(bytes(body), response.headers.get("content-type", ""))

    @staticmethod
    def _parse_response(body: bytes, content_type: str) -> Transcript:
        if not body:
            raise ProviderError("ASR endpoint returned an empty response")
        text: str | None = None
        language: str | None = None
        confidence: float | None = None
        try:
            payload = json.loads(body)
        except (json.JSONDecodeError, UnicodeDecodeError):
            if content_type.lower().startswith("text/plain"):
                text = body.decode("utf-8", errors="strict").strip()
            else:
                raise ProviderError("ASR endpoint returned invalid JSON") from None
        else:
            if isinstance(payload, dict):
                candidate = payload.get("text")
                if isinstance(candidate, str):
                    text = candidate.strip()
                if isinstance(payload.get("language"), str):
                    language = payload["language"][:24]
                candidate_confidence = payload.get("confidence")
                if isinstance(candidate_confidence, (int, float)):
                    value = float(candidate_confidence)
                    if 0 <= value <= 1:
                        confidence = value
        if not text:
            raise ProviderError("ASR endpoint returned no transcript text")
        return Transcript(text=text, language=language, confidence=confidence)

    async def close(self) -> None:
        if self._owns_client:
            await self._client.aclose()


class OpenAICompatibleProvider(_LocalMultipartProvider):
    """Final transcription against an already-running local OpenAI-style endpoint."""

    name = "openai"


class WhisperCppProvider(_LocalMultipartProvider):
    """Final transcription against an already-running whisper.cpp `/inference`."""

    name = "whisper_cpp"

    def form_fields(self, start: ASRStart) -> dict[str, str]:
        fields = super().form_fields(start)
        fields["temperature"] = "0"
        return fields


class TransformersWhisperProvider:
    """In-process, local-files-only Whisper for the MI300X side tap.

    PyTorch intentionally keeps its CUDA-compatible API name on ROCm, so an
    AMD accelerator is selected as ``cuda:0``.  The provider never accepts a
    Hub ID or URL and never enables remote code.
    """

    name = "transformers_local"
    supports_partials = False

    def __init__(
        self,
        *,
        model_path: Path,
        device: str,
        dtype: str = "auto",
        pipeline_runner=None,
    ) -> None:
        self.model_path = Path(model_path)
        self.device = device
        self.dtype = dtype
        self.sample_rate_hz = 16_000
        self._state_lock = asyncio.Lock()
        self._inflight: asyncio.Task | None = None
        self._closed = False
        self._runner = pipeline_runner or self._load_pipeline()

    def _torch_dtype(self, torch):
        if self.device == "cpu":
            if self.dtype not in {"auto", "float32"}:
                raise ConfigurationError(
                    "CPU transformers_local ASR cannot use reduced precision"
                )
            return torch.float32
        selected = "float16" if self.dtype == "auto" else self.dtype
        if selected not in {"float16", "bfloat16", "float32"}:
            raise ConfigurationError("unsupported local Whisper dtype")
        if selected == "bfloat16":
            check_bf16 = getattr(torch.cuda, "is_bf16_supported", None)
            if not callable(check_bf16) or not check_bf16():
                raise ConfigurationError(
                    "OMNI_ASR_DTYPE=bfloat16 is not supported by this PyTorch device"
                )
        return {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
        }[selected]

    def _load_pipeline(self):
        try:
            import torch
            from transformers import (
                AutoModelForSpeechSeq2Seq,
                AutoProcessor,
                pipeline,
            )
        except ImportError as exc:
            raise ConfigurationError(
                "transformers_local requires local torch and transformers packages"
            ) from exc

        if self.device != "cpu":
            prefix, separator, index = self.device.partition(":")
            if prefix != "cuda" or not separator or not index.isdigit():
                raise ConfigurationError("local Whisper device must be cpu or cuda:<index>")
            if not torch.cuda.is_available():
                raise ConfigurationError(
                    "OMNI_ASR_DEVICE requests an accelerator, but PyTorch reports none"
                )
            device_index = int(index)
            if device_index >= torch.cuda.device_count():
                raise ConfigurationError("OMNI_ASR_DEVICE selects a missing accelerator")
        dtype = self._torch_dtype(torch)
        try:
            model = AutoModelForSpeechSeq2Seq.from_pretrained(
                str(self.model_path),
                torch_dtype=dtype,
                low_cpu_mem_usage=True,
                use_safetensors=True,
                local_files_only=True,
                trust_remote_code=False,
                attn_implementation="sdpa",
            )
            processor = AutoProcessor.from_pretrained(
                str(self.model_path),
                local_files_only=True,
                trust_remote_code=False,
            )
            if getattr(model.config, "model_type", None) != "whisper":
                raise ValueError("the local checkpoint is not a Whisper model")
            if not hasattr(processor, "tokenizer") or not hasattr(
                processor, "feature_extractor"
            ):
                raise ValueError("the local Whisper processor is incomplete")
            self.sample_rate_hz = int(processor.feature_extractor.sampling_rate)
            model.to(self.device)
            return pipeline(
                "automatic-speech-recognition",
                model=model,
                tokenizer=processor.tokenizer,
                feature_extractor=processor.feature_extractor,
                torch_dtype=dtype,
                device=self.device,
            )
        except (AttributeError, OSError, RuntimeError, TypeError, ValueError) as exc:
            raise ConfigurationError(
                "the local Whisper checkpoint could not be loaded without downloads"
            ) from exc

    async def partial(self, pcm: bytes, start: ASRStart) -> None:
        return None

    async def final(self, pcm: bytes, start: ASRStart) -> Transcript:
        if not pcm:
            raise ProviderError("No audio received")
        samples = np.frombuffer(pcm, dtype="<i2")
        if samples.size % start.channels:
            raise ProviderError("PCM payload is not aligned to complete sample frames")
        waveform = samples.astype(np.float32) / 32768.0
        if start.channels > 1:
            waveform = waveform.reshape(-1, start.channels).mean(axis=1)
        generate_kwargs: dict[str, str] = {"task": "transcribe"}
        if start.language:
            # Browser locales such as de-DE are not valid Whisper language IDs.
            generate_kwargs["language"] = start.language.replace("_", "-").split("-", 1)[0].lower()

        runner_kwargs: dict[str, object] = {
            "return_timestamps": False,
            "generate_kwargs": generate_kwargs,
        }
        duration_s = waveform.size / start.sample_rate_hz
        if duration_s > 29.0:
            # v4.51's ASR pipeline explicitly accepts these call parameters and
            # stitches overlapping Whisper chunks without retaining audio.
            runner_kwargs["chunk_length_s"] = 30.0
            runner_kwargs["stride_length_s"] = (4.0, 2.0)

        async with self._state_lock:
            if self._closed:
                raise ProviderError("local Whisper provider is closed")
            if self._inflight is not None and not self._inflight.done():
                # asyncio cancellation cannot stop a running GPU worker thread.
                # Rejecting here prevents a timed-out request from overlapping it.
                raise ProviderError("local Whisper is still finishing the previous request")
            runner = self._runner
            if runner is None:
                raise ProviderError("local Whisper provider is closed")

            def transcribe():
                return runner(
                    {"raw": waveform, "sampling_rate": start.sample_rate_hz},
                    **runner_kwargs,
                )

            task = asyncio.create_task(asyncio.to_thread(transcribe))
            self._inflight = task

            def consume_background_result(done: asyncio.Task) -> None:
                try:
                    done.exception()
                except asyncio.CancelledError:
                    pass
                if self._closed:
                    self._runner = None

            task.add_done_callback(consume_background_result)

        try:
            # Shield is essential: wait_for() may cancel this coroutine, but it
            # must not mark the underlying thread future complete prematurely.
            result = await asyncio.shield(task)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            raise ProviderError("local Whisper transcription failed") from exc
        finally:
            if task.done():
                async with self._state_lock:
                    if self._inflight is task:
                        self._inflight = None
        text = result.get("text") if isinstance(result, dict) else None
        if not isinstance(text, str) or not text.strip():
            raise ProviderError("local Whisper returned no transcript text")
        return Transcript(text=text.strip(), language=start.language)

    async def close(self) -> None:
        async with self._state_lock:
            self._closed = True
            if self._inflight is None or self._inflight.done():
                self._runner = None


def build_provider(settings: ASRSettings) -> ASRProvider:
    settings.validated()
    if settings.provider == "mock":
        return MockASRProvider()
    if settings.provider == "transformers_local":
        if not settings.model:
            raise ConfigurationError("A local ASR model directory is required")
        return TransformersWhisperProvider(
            model_path=Path(settings.model),
            device=settings.device,
            dtype=settings.dtype,
        )
    if not settings.endpoint or not settings.model:
        # Kept here as a second fail-closed guard for programmatic callers.
        raise ConfigurationError("A local ASR endpoint and model are required")
    provider_type = (
        OpenAICompatibleProvider if settings.provider == "openai" else WhisperCppProvider
    )
    return provider_type(
        endpoint=settings.endpoint,
        model=settings.model,
        api_key=settings.api_key,
        connect_timeout_s=settings.connect_timeout_s,
        write_timeout_s=settings.write_timeout_s,
        transcribe_timeout_s=settings.transcribe_timeout_s,
        max_response_bytes=settings.max_response_bytes,
    )