"""Backend-neutral transcription and realtime turn abstractions."""

from __future__ import annotations

from dataclasses import dataclass
import json
from queue import Empty, Full, Queue
import threading
import time
from typing import Any, Protocol, runtime_checkable

from .recorder import Capture

MAX_PCM_QUEUE_BYTES = 240_000  # five seconds of 24 kHz mono PCM16
MAX_PCM_QUEUE_CHUNKS = 64
MAX_PROVIDER_EVENT_BYTES = 64 * 1024
MAX_PROVIDER_EVENT_QUEUE_BYTES = 512 * 1024
MAX_PROVIDER_EVENT_QUEUE_ITEMS = 128
MAX_PROVIDER_TEXT_BYTES = 32 * 1024
RECEIVE_TIMEOUT = object()


class ProviderError(RuntimeError):
    """A bounded, redacted provider failure suitable for the service boundary."""

    def __init__(self, code: str, detail: str = "transcription provider failed") -> None:
        self.code = code
        try:
            self.detail = detail.encode("utf-8", errors="replace")[:2048].decode("utf-8", errors="ignore")
        except (AttributeError, UnicodeError):
            self.detail = "transcription provider failed"
        super().__init__(self.detail)


@dataclass(frozen=True)
class BackendDescriptor:
    backend: str
    model: str


@runtime_checkable
class TranscriptionTurn(Protocol):
    sample_rate_hz: int
    retain_pcm: bool

    def enqueue_pcm(self, chunk: bytes) -> None: ...
    def finish(self, capture: Capture | None = None) -> str: ...
    def poll_error(self) -> ProviderError | None: ...
    def cancel(self) -> None: ...
    def close(self) -> None: ...

    @property
    def cleanup_complete(self) -> bool: ...


@runtime_checkable
class TranscriptionBackend(Protocol):
    descriptor: BackendDescriptor
    sample_rate_hz: int

    def start_turn(self) -> TranscriptionTurn: ...

    def cancel_prepare(self) -> None: ...


class PcmQueue:
    """A byte- and item-bounded nonblocking queue for recorder drain threads."""

    def __init__(
        self,
        *,
        max_bytes: int = MAX_PCM_QUEUE_BYTES,
        max_chunks: int = MAX_PCM_QUEUE_CHUNKS,
    ) -> None:
        if max_bytes <= 0 or max_chunks <= 0:
            raise ValueError("PCM queue bounds must be positive")
        self._queue: Queue[bytes | None] = Queue(maxsize=max_chunks)
        self._max_bytes = max_bytes
        self._queued_bytes = 0
        self._lock = threading.Lock()
        self._close_lock = threading.Lock()
        self._closed = False
        self._end_queued = False

    @property
    def queued_bytes(self) -> int:
        with self._lock:
            return self._queued_bytes

    def put(self, chunk: bytes) -> None:
        if not isinstance(chunk, bytes) or not chunk:
            raise ProviderError("queue_invalid", "provider received invalid audio")
        with self._lock:
            if self._closed:
                raise ProviderError("queue_closed", "provider audio queue is closed")
            if len(chunk) > self._max_bytes or self._queued_bytes + len(chunk) > self._max_bytes:
                raise ProviderError("queue_overflow", "provider audio queue is full")
            self._queued_bytes += len(chunk)
        try:
            self._queue.put_nowait(chunk)
        except Full as error:
            with self._lock:
                self._queued_bytes -= len(chunk)
            raise ProviderError("queue_overflow", "provider audio queue is full") from error

    def close(self, *, timeout: float = 15.0) -> None:
        if not self._close_lock.acquire(timeout=max(0.0, timeout)):
            raise ProviderError("queue_overflow", "provider audio queue could not be finalized")
        try:
            with self._lock:
                if self._end_queued:
                    return
                self._closed = True
            # A full queue still has valid audio that must be drained. The
            # closed flag is a side-channel EOF marker when the sentinel cannot
            # be inserted yet, rather than dropping the queued tail.
            try:
                self._queue.put_nowait(None)
            except Full:
                pass
            with self._lock:
                self._end_queued = True
        finally:
            self._close_lock.release()

    def get(self, timeout: float | None = None) -> bytes | None:
        deadline = None if timeout is None else time.monotonic() + timeout
        while True:
            poll_timeout = 0.1 if deadline is None else max(0.0, min(0.1, deadline - time.monotonic()))
            try:
                chunk = self._queue.get(timeout=poll_timeout)
            except Empty:
                with self._lock:
                    closed = self._closed
                if closed:
                    return None
                if deadline is not None and time.monotonic() >= deadline:
                    raise
                continue
            if chunk is not None:
                with self._lock:
                    self._queued_bytes -= len(chunk)
            return chunk


class LocalNemoTurn:
    """Post-stop local inference turn retaining audio only in the recorder."""

    sample_rate_hz = 16_000
    retain_pcm = True

    def __init__(self, model: Any) -> None:
        self._model = model
        self._cancelled = False
        self._closed = False
        self._finished = False

    @property
    def cleanup_complete(self) -> bool:
        return self._closed

    def enqueue_pcm(self, chunk: bytes) -> None:
        # Local turns use Capture.pcm so no second copy is made here.
        if not self.retain_pcm:
            raise ProviderError("local_invalid", "local turn cannot stream audio")

    def finish(self, capture: Capture | None = None) -> str:
        if self._cancelled:
            raise ProviderError("local_cancelled", "local turn was cancelled")
        if self._closed or self._finished:
            raise ProviderError("local_invalid", "local turn was already finalized")
        self._finished = True
        if capture is None:
            raise ProviderError("local_invalid", "local capture is missing")
        if capture.sample_rate_hz != self.sample_rate_hz:
            raise ProviderError("local_invalid", "local capture has an unsupported sample rate")
        from .audio import pcm_s16le_to_floats

        max_bytes = getattr(self._model, "max_pcm_bytes", 5 * 60 * 16_000 * 2)
        samples = pcm_s16le_to_floats(capture.pcm, max_bytes=max_bytes)
        result = self._model.transcribe(samples)
        if not isinstance(result, str):
            raise ProviderError("local_failure", "local model returned an invalid transcript")
        return result

    def poll_error(self) -> ProviderError | None:
        return None

    def cancel(self) -> None:
        self._cancelled = True

    def close(self) -> None:
        self._closed = True


class LocalNemoBackend:
    """Backend wrapper preserving the lazy, warm NeMo model behavior."""

    sample_rate_hz = 16_000

    def __init__(self, model: Any) -> None:
        self._model = model
        self.descriptor = BackendDescriptor("local", "nemo-parakeet")

    @property
    def ready(self) -> bool:
        return bool(getattr(self._model, "loaded", False))

    def start_turn(self) -> LocalNemoTurn:
        self._model.load()
        return LocalNemoTurn(self._model)

    def cancel_prepare(self) -> None:
        # NeMo model loading is not interruptible, but the service tombstone
        # prevents a cancelled preparation from starting capture.
        return None


class QueuedCloudTurn:
    """Common ordered sender/receiver for cloud adapters.

    Provider handshakes happen before this object is created. The recorder only
    calls ``enqueue_pcm``; it never performs network I/O. A sender thread drains
    the bounded queue and sends the end marker after every queued PCM chunk. A
    receiver thread continuously consumes provider events during capture so a
    websocket's receive buffer cannot stall the sender.
    """

    sample_rate_hz = 24_000
    retain_pcm = False

    def __init__(
        self,
        transport: Any,
        *,
        timeout: float = 15.0,
        max_queue_bytes: int = MAX_PCM_QUEUE_BYTES,
        max_queue_chunks: int = MAX_PCM_QUEUE_CHUNKS,
    ) -> None:
        if max_queue_bytes <= 0 or max_queue_bytes > MAX_PCM_QUEUE_BYTES:
            raise ValueError("cloud PCM backlog must not exceed five seconds")
        self._transport = transport
        self._timeout = timeout
        self._queue = PcmQueue(max_bytes=max_queue_bytes, max_chunks=max_queue_chunks)
        self._provider_events: Queue[tuple[dict[str, Any], int]] = Queue(maxsize=MAX_PROVIDER_EVENT_QUEUE_ITEMS)
        self._provider_event_bytes = 0
        self._provider_closed = False
        self._sender_error: ProviderError | None = None
        self._receiver_error: ProviderError | None = None
        self._event_lock = threading.Lock()
        self._cancel_event = threading.Event()
        self._transport_close_lock = threading.Lock()
        self._transport_close_thread: threading.Thread | None = None
        self._cancelled = False
        self._finished = False
        self._closed = False
        self._sent_bytes = 0
        self._stream_started_at = time.monotonic()
        self._sender = threading.Thread(target=self._send_loop, daemon=True, name="dictate-provider-send")
        self._receiver = threading.Thread(target=self._receive_loop, daemon=True, name="dictate-provider-recv")
        self._sender.start()
        self._receiver.start()

    @property
    def cleanup_complete(self) -> bool:
        close_thread = self._transport_close_thread
        return (
            not self._sender.is_alive()
            and not self._receiver.is_alive()
            and (close_thread is None or not close_thread.is_alive())
        )

    def enqueue_pcm(self, chunk: bytes) -> None:
        if self._cancelled or self._finished:
            raise ProviderError(f"{self._provider_prefix}_queue_closed", "provider turn is no longer accepting audio")
        try:
            self._queue.put(chunk)
        except ProviderError as error:
            raise ProviderError(f"{self._provider_prefix}_{error.code}", error.detail) from error

    def close_input(self) -> None:
        """Signal EOF after the recorder has drained stdout to EOF."""
        self._queue.close()

    def poll_error(self) -> ProviderError | None:
        error = self._sender_error or self._receiver_error
        if error is not None:
            return error
        with self._event_lock:
            provider_closed = self._provider_closed
        if provider_closed and not self._finished and not self._cancelled:
            return ProviderError(
                self._provider_prefix + "_premature_close",
                "provider closed before recording was finalized",
            )
        return None

    def finish(self, capture: Capture | None = None) -> str:
        if self._cancelled:
            raise ProviderError("cancelled", "provider turn was cancelled")
        if self._finished:
            raise ProviderError("invalid_state", "provider turn was already finalized")
        self._finished = True
        try:
            self._queue.close(timeout=self._timeout)
        except ProviderError:
            self.cancel()
            raise
        self._sender.join(timeout=self._timeout)
        if self._sender.is_alive():
            self.cancel()
            raise ProviderError(self._provider_prefix + "_timeout", "provider audio sender timed out")
        if self._sender_error is not None:
            error = self._sender_error
            self.close()
            raise error
        try:
            transcript = self._finish_provider()
            # A correlated final isn't enough to publish a result while the
            # receiver can still report a queued provider failure. Stop and
            # join it first, then make the receiver error check authoritative.
            self._cancel_event.set()
            self._request_transport_close()
            self._receiver.join(timeout=self._timeout)
            if self._receiver.is_alive():
                raise ProviderError(
                    self._provider_prefix + "_timeout",
                    "provider receiver did not stop after the final transcript",
                )
            if self._receiver_error is not None:
                raise self._receiver_error
            return transcript
        finally:
            self.close()

    def cancel(self) -> None:
        if self._cancelled and self.cleanup_complete:
            return
        self._cancelled = True
        self._cancel_event.set()
        self._request_transport_close()
        try:
            self._queue.close(timeout=0.1)
        except ProviderError:
            pass
        # Closing the transport is the cancellation signal. Do not block the
        # protocol loop waiting for provider threads; ``cleanup_complete``
        # keeps the tombstoned generation busy until they actually exit.

    def close(self) -> None:
        if self._closed and self.cleanup_complete:
            return
        self._cancel_event.set()
        self._request_transport_close()
        # Keep close bounded and nonblocking from the service loop. A later
        # poll retries the joins while cleanup_complete prevents a new turn.
        for thread in (self._sender, self._receiver, self._transport_close_thread):
            if thread is not None:
                thread.join(timeout=min(self._timeout, 0.05))
        self._closed = self.cleanup_complete

    def _send_loop(self) -> None:
        try:
            while True:
                chunk = self._queue.get(timeout=self._timeout)
                if chunk is None:
                    # Queue closure means commit only for a successful recorder
                    # EOF. Cancellation closes the same bounded queue merely to
                    # wake this worker and must never finalize partial audio.
                    if not self._cancel_event.is_set():
                        self._send_end()
                    return
                # Pace the stream against the declared PCM rate. This prevents
                # a fast recorder from building more than five seconds of
                # provider-side backlog while preserving ordered sends.
                target_elapsed = self._sent_bytes / (self.sample_rate_hz * 2)
                elapsed = time.monotonic() - self._stream_started_at
                if target_elapsed > elapsed and self._cancel_event.wait(target_elapsed - elapsed):
                    return
                if self._cancel_event.is_set():
                    return
                self._send_audio(chunk)
                self._sent_bytes += len(chunk)
        except ProviderError as error:
            self._sender_error = error
            self._abort_transport()
        except Empty:
            self._sender_error = ProviderError(
                self._provider_prefix + "_timeout", "provider audio sender timed out"
            )
            self._abort_transport()
        except Exception:
            self._sender_error = ProviderError(
                self._provider_prefix + "_send", "provider audio upload failed"
            )
            self._abort_transport()

    def _receive_loop(self) -> None:
        try:
            while not self._cancel_event.is_set():
                event = self._receive_provider_event(timeout=0.1)
                if event is RECEIVE_TIMEOUT:
                    continue
                if event is None:
                    with self._event_lock:
                        self._provider_closed = True
                    return
                if not isinstance(event, dict):
                    raise ProviderError(self._provider_prefix + "_malformed_event", "provider event is invalid")
                try:
                    size = len(json.dumps(event, ensure_ascii=False, separators=(",", ":")).encode("utf-8"))
                except (TypeError, UnicodeError, ValueError, RecursionError, OverflowError) as error:
                    raise ProviderError(self._provider_prefix + "_malformed_event", "provider event is invalid") from error
                with self._event_lock:
                    if self._provider_event_bytes + size > MAX_PROVIDER_EVENT_QUEUE_BYTES:
                        raise ProviderError(self._provider_prefix + "_event_overflow", "provider event queue is full")
                    self._provider_event_bytes += size
                try:
                    self._provider_events.put_nowait((event, size))
                except Full as error:
                    with self._event_lock:
                        self._provider_event_bytes -= size
                    raise ProviderError(self._provider_prefix + "_event_overflow", "provider event queue is full") from error
        except ProviderError as error:
            # Preserve a provider failure that races with successful
            # finalization. Cancellation does not consume this error, but it
            # avoids recursively closing a transport we are already closing.
            self._receiver_error = error
            if not self._cancel_event.is_set():
                self._abort_transport()
        except Exception:
            if not self._cancel_event.is_set():
                self._receiver_error = ProviderError(
                    self._provider_prefix + "_receive", "provider event stream failed"
                )
                self._abort_transport()

    def _next_provider_event(self, timeout: float) -> dict[str, Any] | None:
        try:
            event, size = self._provider_events.get(timeout=max(0.0, timeout))
        except Empty as error:
            error_from_receiver = self._receiver_error
            if error_from_receiver is not None:
                raise error_from_receiver
            with self._event_lock:
                closed = self._provider_closed
            if closed:
                return None
            raise ProviderError(self._provider_prefix + "_timeout", "timed out waiting for provider event") from error
        with self._event_lock:
            self._provider_event_bytes -= size
        return event

    @property
    def _provider_prefix(self) -> str:
        return "provider"

    def _send_audio(self, chunk: bytes) -> None:
        raise NotImplementedError

    def _send_end(self) -> None:
        raise NotImplementedError

    def _finish_provider(self) -> str:
        raise NotImplementedError

    def _receive_provider_event(self, *, timeout: float) -> dict[str, Any] | None | object:
        raise NotImplementedError

    def _request_transport_close(self) -> None:
        with self._transport_close_lock:
            if self._transport_close_thread is not None:
                return
            self._transport_close_thread = threading.Thread(
                target=self._abort_transport,
                daemon=True,
                name="dictate-provider-close",
            )
            self._transport_close_thread.start()

    def _abort_transport(self) -> None:
        try:
            self._transport.close()
        except Exception:
            pass


def bounded_provider_text(value: Any, *, field: str = "transcript") -> str:
    if not isinstance(value, str):
        raise ProviderError("provider_malformed", f"provider {field} is invalid")
    try:
        if len(value.encode("utf-8")) > MAX_PROVIDER_TEXT_BYTES:
            raise ProviderError("provider_malformed", f"provider {field} is too large")
    except (UnicodeEncodeError, AttributeError) as error:
        raise ProviderError("provider_malformed", f"provider {field} is invalid") from error
    return value


def provider_event(value: Any) -> dict[str, Any]:
    """Validate one already decoded event and apply a serialized-size bound."""

    if not isinstance(value, dict) or any(not isinstance(key, str) for key in value):
        raise ProviderError("provider_malformed", "provider event is not an object")
    try:
        encoded = json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    except (TypeError, UnicodeError, ValueError, RecursionError, OverflowError) as error:
        raise ProviderError("provider_malformed", "provider event is invalid") from error
    if len(encoded) > MAX_PROVIDER_EVENT_BYTES:
        raise ProviderError("provider_malformed", "provider event is too large")
    return value


def backend_from_config(config: Any) -> TranscriptionBackend:
    """Construct only the selected backend, keeping cloud imports lazy."""
    if config.backend == "local":
        from .backends.local import LocalNemoBackend
        return LocalNemoBackend(config.local)
    if config.backend == "meta":
        from .backends.meta import MetaBackend
        return MetaBackend(config.meta)
    if config.backend == "openai":
        from .backends.openai import OpenAIBackend
        return OpenAIBackend(config.openai)
    raise ProviderError("configuration", "unsupported transcription backend")


__all__ = [
    "BackendDescriptor", "LocalNemoBackend", "LocalNemoTurn", "MAX_PCM_QUEUE_BYTES",
    "MAX_PROVIDER_EVENT_BYTES", "MAX_PROVIDER_EVENT_QUEUE_BYTES", "MAX_PROVIDER_TEXT_BYTES",
    "PcmQueue", "ProviderError", "QueuedCloudTurn", "RECEIVE_TIMEOUT",
    "TranscriptionBackend", "TranscriptionTurn", "backend_from_config",
    "bounded_provider_text", "provider_event",
]