"""Speech: a microphone's audio to text, and text to audio a speaker plays.

Both talk to an OpenAI-compatible speech server, such as `parrot serve` kept
running by a Service node, whose address arrives on a ``url`` input. An empty
url is a server that is not answering: the node reads degraded and drops what
it gets.
"""

from __future__ import annotations

import asyncio
import contextlib
import hashlib
import logging
import re
import threading
import time
from collections.abc import Coroutine, Iterable
from concurrent.futures import Future
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel, ConfigDict, Field

from fluksio.ai import client
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.base import Node

if TYPE_CHECKING:
    from fastapi import FastAPI

logger = logging.getLogger(__name__)

#: Seconds a session waits for its next chunk, and the most one utterance may
#: take from its first chunk to its text. Past either it is given up.
IDLE_S = 10.0
UTTERANCE_S = 60.0
#: The rate of a microphone's PCM chunks, the panels' (openpanel V3).
RATE = 16000
#: How many spoken sentences a tts node keeps.
SPOKEN = 256
#: How long after its estimated end a sentence still counts as playing: the
#: device fetches it before it starts.
PLAYING_SLACK_S = 1.0
#: The partial for an utterance that came back empty, often a wake word that
#: was not one, or was lost: shown on the caption, never spoken.
NOTHING_HEARD = "Didn't catch that."
#: Units read out after a number, the longer first. Kokoro says "22 °C" as
#: "22 degrees C" and "3.4 kWh" letter by letter.
UNITS = {
    "°C": "degrees",
    "°": "degrees",
    "kWh": "kilowatt hours",
    "kW": "kilowatts",
    "W": "watts",
}


def _check(specs: list[MessageSpec], ports: dict[str, DType], node: str) -> None:
    """Refuse a port the node does not have, or one of another type."""
    for spec in specs:
        dtype = ports.get(spec.port)
        if dtype is None:
            raise ValueError(
                f"'{spec.port}' is not a port of {node}. Its ports are "
                f"{', '.join(ports)}."
            )
        if spec.dtype is not dtype:
            raise ValueError(f"{node}'s '{spec.port}' is {dtype.value}.")


def _is_pcm(audio: dict[str, Any]) -> bool:
    """A microphone chunk rather than a whole clip."""
    media_type = str(audio.get("media_type") or "").lower()
    return media_type.startswith(("audio/l16", "audio/pcm"))


def sentences(text: str) -> list[str]:
    """``text`` cut after each ``.``, ``!`` and ``?``, one call to speak each."""
    return [s for s in (p.strip() for p in re.split(r"(?<=[.!?])\s+", text)) if s]


def speakable(text: str) -> str:
    """``text`` as it is said: a unit after a number in words, and a clock
    time as its hour and minutes, since "07:30" is read "zero seven thirty"."""
    units = "|".join(map(re.escape, UNITS))
    text = re.sub(
        rf"(\d)\s*({units})(?![A-Za-z])", lambda m: f"{m[1]} {UNITS[m[2]]}", text
    )
    return re.sub(r"\b([01]?\d|2[0-3]):([0-5]\d)\b", _clock, text)


def _clock(time: re.Match[str]) -> str:
    hour, minute = int(time[1]), int(time[2])
    if minute == 0:
        return f"{hour} o'clock"
    return f"{hour} oh {minute}" if minute < 10 else f"{hour} {minute}"


def mp3_seconds(data: bytes) -> float:
    """An MP3's length, from the frame count LAME writes into its first frame."""
    header = int.from_bytes(data[:4], "big")
    version = (header >> 19) & 3  # 3 is MPEG-1, 2 MPEG-2, 0 MPEG-2.5
    rates = {
        3: (44100, 48000, 32000),
        2: (22050, 24000, 16000),
        0: (11025, 12000, 8000),
    }
    rate = rates[version][(header >> 10) & 3]
    at = max(data.find(b"Xing", 0, 64), data.find(b"Info", 0, 64))
    if at < 0:
        raise ValueError("the MP3 has no Xing or Info header to read its length from")
    frames = int.from_bytes(data[at + 8 : at + 12], "big")
    return round(frames * (1152 if version == 3 else 576) / rate, 3)


class _Speech(Node):
    """What both nodes share: the url, the artifact store and their health."""

    __slots__ = ("_artifacts", "_loop")

    def __init__(
        self,
        f: Any,
        requires: Iterable[MessageSpec],
        provides: Iterable[MessageSpec],
        params: dict[str, Any] | None,
        name: str,
        inputs: dict[str, DType],
        outputs: dict[str, DType],
    ):
        # Set on start; a test sets them directly.
        self._artifacts: Any = None
        self._loop: asyncio.AbstractEventLoop | None = None
        super().__init__(
            f, requires=requires, provides=provides, params=params, name=name
        )
        _check(self.input_ports, inputs, name)
        _check(self.output_ports, outputs, name)
        if not any(spec.port == "url" for spec in self.input_ports):
            raise ValueError(
                f"{name} needs a 'url' input: the speech server's address."
            )

    async def start(self, app: FastAPI | None = None) -> None:
        self._loop = asyncio.get_running_loop()
        self._artifacts = getattr(app.state, "artifact_store", None) if app else None

    def _reachable(self, url: str) -> bool:
        """Whether there is a server to ask, said as the node's health."""
        if not url:
            if self._health != "degraded":
                self.report_health("degraded", "no speech server: its url is empty")
            return False
        if self._health != "ok":
            self.report_health("ok")
        return True

    def _load(self, reference: dict[str, Any]) -> bytes:
        return b"".join(self._artifacts.read(str(reference["digest"])))

    def _call(self, work: Coroutine[Any, Any, Any]) -> Future[Any]:
        """Run on the engine's loop, where the server's connection is kept."""
        if self._loop is None:
            work.close()
            raise RuntimeError(f"{self.name} only runs in a live flow")
        return asyncio.run_coroutine_threadsafe(work, self._loop)

    def _say(self, **values: Any) -> None:
        """Publish on the ports that are wired."""
        wired = {spec.port for spec in self.output_ports if spec.name}
        self.emit({port: value for port, value in values.items() if port in wired})


class SttNode(_Speech):
    """Speech to text: a microphone's audio in, what was said out.

    Inputs: ``audio``, a stream of 16 kHz PCM16 chunks or a whole clip,
    ``listening``, true while the device listens, and ``url``.
    Outputs: ``partial``, the words agreed on so far, and ``text``.

    One realtime session per utterance, opened when ``listening`` turns true
    or at the first chunk, and committed when it turns false. A clip goes to
    the server whole. Either way the node returns at once, since a device's
    next frame waits for it, and publishes once the server answers. Nothing
    heard, or an utterance lost and logged, is a short note on ``partial``
    alone.
    """

    __slots__ = ("_heard", "_listening", "_lock", "_session", "_tasks")

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        self._lock = threading.Lock()
        #: The chunks of the utterance being spoken, None between utterances.
        self._session: asyncio.Queue[bytes | None] | None = None
        #: The digest last heard, and whether the device was listening.
        self._heard = ""
        self._listening: bool | None = None
        self._tasks: set[Future[None]] = set()
        inputs = [
            # The url is read, not waited for: its changing is nothing to hear.
            spec.model_copy(update={"trigger": spec.port != "url"})
            for spec in self._normalize_ports(requires)
        ]
        if not any(spec.port == "audio" for spec in inputs):
            raise ValueError("stt needs an 'audio' input: the microphone.")
        super().__init__(
            self._hear,
            inputs,
            self._normalize_ports(provides),
            params,
            name or "stt",
            {"audio": DType.AUDIO, "listening": DType.BOOL, "url": DType.STR},
            {"partial": DType.STR, "text": DType.STR},
        )

    async def start(self, app: FastAPI | None = None) -> None:
        await super().start(app)
        if self._pipeline is None:
            return
        # What the inputs hold already was heard before this start.
        held = await asyncio.to_thread(
            self._pipeline.state.get_present, list(self.requires)
        )
        for spec in self.input_ports:
            value = held.get(spec.name)
            if spec.port == "audio" and isinstance(value, dict):
                self._heard = str(value.get("digest") or "")
            elif spec.port == "listening" and isinstance(value, bool):
                self._listening = value

    async def stop(self, app: FastAPI | None = None) -> None:
        for task in list(self._tasks):
            task.cancel()
        self._session = None

    def _hear(
        self,
        params: dict[str, Any],
        audio: dict[str, Any] | None = None,
        listening: bool | None = None,
        url: str = "",
    ) -> None:
        with self._lock:
            digest = str((audio or {}).get("digest") or "")
            fresh = bool(digest) and digest != self._heard
            self._heard = digest or self._heard
            was, self._listening = self._listening, listening
            if not self._reachable(url):
                self._session = None
                return None
            if fresh and audio is not None and not _is_pcm(audio):
                self._spawn(self._clip(url, audio, self._load(audio)))
            elif fresh and audio is not None:
                if self._session is None:
                    self._session = self._open(url)
                self._feed(self._session, self._load(audio))
            elif listening and not was and self._session is None:
                self._session = self._open(url)
            elif was and not listening and self._session is not None:
                self._feed(self._session, None)
                self._session = None
        return None

    def _spawn(self, work: Coroutine[Any, Any, None]) -> None:
        """Run on the engine's loop, so this run can return at once."""
        task = self._call(work)
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)

    def _open(self, url: str) -> asyncio.Queue[bytes | None]:
        chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
        self._spawn(self._transcribe(url, chunks))
        # The reply's first sentence then finds the connection open.
        self._spawn(self._warm(url))
        return chunks

    async def _warm(self, url: str) -> None:
        with contextlib.suppress(Exception):
            await client.warm(url + "/v1")

    def _feed(self, chunks: asyncio.Queue[bytes | None], chunk: bytes | None) -> None:
        assert self._loop is not None
        self._loop.call_soon_threadsafe(chunks.put_nowait, chunk)

    async def _transcribe(self, url: str, chunks: asyncio.Queue[bytes | None]) -> None:
        words = ""
        try:
            async with asyncio.timeout(UTTERANCE_S):
                async for final, text in client.transcribe_stream(
                    url + "/v1", chunks, rate=RATE, idle_s=IDLE_S
                ):
                    words = text.strip() if final else words + text
                    if final and words:
                        await asyncio.to_thread(self._say, partial=words, text=words)
                    elif final:
                        await asyncio.to_thread(self._say, partial=NOTHING_HEARD)
                    elif words:
                        await asyncio.to_thread(self._say, partial=words)
        except Exception as exc:
            logger.warning(
                "'%s' lost an utterance: %s: %s", self.id, type(exc).__name__, exc
            )
            await asyncio.to_thread(self._say, partial=NOTHING_HEARD)
        finally:
            # Ended without a commit, by a timeout or a failure: the next chunk
            # opens a session of its own rather than feeding this dead one.
            with self._lock:
                if self._session is chunks:
                    self._session = None

    async def _clip(self, url: str, audio: dict[str, Any], data: bytes) -> None:
        try:
            text = await client.transcribe(
                url + "/v1", data, "speech.wav", str(audio.get("media_type"))
            )
        except Exception as exc:
            logger.warning("'%s' lost a clip: %s: %s", self.id, type(exc).__name__, exc)
            text = ""
        if text.strip():
            await asyncio.to_thread(self._say, partial=text.strip(), text=text.strip())
        else:
            await asyncio.to_thread(self._say, partial=NOTHING_HEARD)


class TtsNode(_Speech):
    """Text to speech: each sentence of ``text`` spoken, out on ``speech``.

    One call to the server per sentence, each sentence published as an
    ``audio`` reference as soon as it is spoken, so a reply plays before its
    last sentence is. A sentence spoken before in the same voice is not spoken
    again.

    A reference is marked ``"append": true`` while the ones before it are
    still playing, which queues it on the device. The first sentence of a
    reply carries no mark and cuts off whatever was playing, however its
    sentences arrive: as one text, or one value each while they are written.
    """

    class Params(BaseModel):
        model_config = ConfigDict(extra="allow")

        voice: str = Field(
            default="af_heart",
            description="The voice, one of Kokoro's such as af_heart or bm_george.",
        )

    __slots__ = ("cfg", "_lock", "_until")

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        self.cfg = self.Params.model_validate(params or {})
        # One reply at a time, so its sentences go out in order.
        self._lock = threading.Lock()
        #: When what was published so far stops playing, on the monotonic clock.
        self._until = 0.0
        inputs = [
            spec.model_copy(update={"trigger": spec.port != "url"})
            for spec in self._normalize_ports(requires)
        ]
        if not any(spec.port == "text" for spec in inputs):
            raise ValueError("tts needs a 'text' input: what to say.")
        super().__init__(
            self._speak,
            inputs,
            self._normalize_ports(provides),
            params,
            name or "tts",
            {"text": DType.STR, "url": DType.STR},
            {"speech": DType.AUDIO},
        )

    def _speak(self, params: dict[str, Any], text: str = "", url: str = "") -> None:
        if not text.strip() or not self._reachable(url):
            return None
        with self._lock:
            for sentence in sentences(speakable(text)):
                reference = self._spoken(url, sentence)
                now = time.monotonic()
                if now < self._until + PLAYING_SLACK_S:
                    reference = {**reference, "append": True}
                self._until = max(now, self._until) + reference["meta"]["seconds"]
                self._say(speech=reference)
        return None

    def _spoken(self, url: str, sentence: str) -> dict[str, Any]:
        """The reference to ``sentence`` spoken, from the cache or the server."""
        key = hashlib.sha256(f"mp3\n{self.cfg.voice}\n{sentence}".encode()).hexdigest()
        # Kept in state, so the sweep keeps the bytes it refers to.
        spoken = dict(self.recall("spoken") or {})
        reference = spoken.get(key)
        if reference is not None and self._artifacts.has(reference["digest"]):
            return dict(reference)
        data = self._call(client.speak(url + "/v1", sentence, self.cfg.voice)).result()
        stored: dict[str, Any] = self._artifacts.put(
            [data], name="speech.mp3", media_type="audio/mpeg"
        )
        reference = stored | {"meta": {"seconds": mp3_seconds(data)}}
        spoken.pop(key, None)
        spoken[key] = reference
        self.remember("spoken", dict(list(spoken.items())[-SPOKEN:]))
        return reference