"""Speech-to-text for one utterance, through the Bolt server's /desk/stt
websocket relay (to Deepgram) — the route Bolt-Pet and the console's mic use,
authenticated with the desk API key; no Deepgram account here.

    t = Transcription(connect)
    t.add(frames)        # as the utterance is spoken (streamed live)
    text = t.finish(audio)   # when it ends; `audio` is the whole utterance
    t.cancel()           # it was too short to matter

Live when possible, so the text is ready the moment you stop talking. If the
socket never opened or died, finish() opens a fresh one and sends `audio`
in one go, from an empty transcript (so a half-live session can't return
words twice). A healthy socket is never re-sent — that would pay twice.
"""

from __future__ import annotations

import json
import threading
import urllib.parse
from typing import Callable, List, Optional

from . import config


class SttError(Exception):
    pass


def connect(sample_rate: int = config.SAMPLE_RATE):
    import websocket  # websocket-client

    query = urllib.parse.urlencode({"session_id": config.SESSION_ID, "encoding": "linear16",
                                    "sample_rate": sample_rate})
    return websocket.create_connection(
        config.SERVER_URL.replace("http", "ws", 1) + "/desk/stt?" + query,
        header=[f"X-Desk-Api-Key: {config.API_KEY}"], timeout=10)


def parse_result(message) -> tuple:
    """(text, is_final) from a Deepgram results frame; ("", False) otherwise."""
    try:
        if isinstance(message, bytes):
            message = message.decode("utf-8", "ignore")
        data = json.loads(message)
    except (TypeError, ValueError):
        return "", False
    if not isinstance(data, dict) or data.get("type") not in (None, "Results"):
        return "", False
    alternatives = (data.get("channel") or {}).get("alternatives") or []
    if not alternatives:
        return "", False
    text = str((alternatives[0] or {}).get("transcript") or "").strip()
    return text, bool(data.get("is_final") or data.get("speech_final"))


class Transcription:
    def __init__(self, connect_fn: Callable = connect):
        self._connect = connect_fn
        self._finals: List[str] = []
        self._interim = ""
        self._lock = threading.Lock()
        self._socket = None
        self._ok = False
        self._reader: Optional[threading.Thread] = None
        self._open()

    def _open(self) -> None:
        try:
            self._socket = self._connect(config.SAMPLE_RATE)
            self._ok = True
        except Exception:
            self._socket, self._ok = None, False
            return
        self._reader = threading.Thread(target=self._read, name="stt-reader", daemon=True)
        self._reader.start()

    def add(self, frames: List[bytes]) -> None:
        for frame in frames:
            self._send(frame)

    def _send(self, data: bytes) -> None:
        if self._ok:
            try:
                self._socket.send_binary(data)
            except Exception:
                self._ok = False

    def finish(self, audio: bytes) -> str:
        """The transcript. Raises SttError only if even the fallback couldn't
        reach the server."""
        if self._ok:
            self._flush()
            return self._text()
        self._close()
        if self._reader is not None:
            self._reader.join(timeout=1)
        with self._lock:
            self._finals, self._interim = [], ""
        self._open()
        if not self._ok:
            raise SttError("couldn't reach the transcription server")
        step = config.FRAME_SAMPLES * 2 * 8
        for offset in range(0, len(audio), step):
            self._send(audio[offset:offset + step])
        self._flush()
        return self._text()

    def cancel(self) -> None:
        self._ok = False
        self._close()

    def _flush(self) -> None:
        """Deepgram sends its last results after CloseStream; wait for them,
        bounded, so a hung socket can't hold up the reply."""
        try:
            self._socket.send(json.dumps({"type": "CloseStream"}))
        except Exception:
            pass
        if self._reader is not None:
            self._reader.join(timeout=4)
        self._close()

    def _close(self) -> None:
        socket, self._socket = self._socket, None
        if socket is not None:
            try:
                socket.close()
            except Exception:
                pass

    def _read(self) -> None:
        socket = self._socket
        while socket is not None:
            try:
                message = socket.recv()
            except Exception:
                break
            if not message:
                break
            text, is_final = parse_result(message)
            with self._lock:
                if is_final:
                    if text:
                        self._finals.append(text)
                    self._interim = ""
                elif text:
                    self._interim = text

    def _text(self) -> str:
        with self._lock:
            return " ".join(self._finals + ([self._interim] if self._interim else [])).strip()