"""The speech nodes: one session per utterance, and each sentence spoken once."""

import asyncio
import hashlib
import io
import wave
from typing import Any

from fluksio.ai import client
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.speech import SttNode, TtsNode, sentences
from fluksio.flow.pipeline import Pipeline

URL = "http://speech:8100"


class Store:
    """The artifact store's surface the nodes use."""

    def __init__(self) -> None:
        self.data: dict[str, bytes] = {}

    def put(self, chunks: list[bytes], name: str, media_type: str) -> dict[str, Any]:
        data = b"".join(chunks)
        digest = "sha256:" + hashlib.sha256(data).hexdigest()
        self.data[digest] = data
        return {"digest": digest, "size": len(data), "media_type": media_type}

    def has(self, digest: str) -> bool:
        return digest in self.data

    def read(self, digest: str) -> list[bytes]:
        return [self.data[digest]]


def test_sentences_are_cut_after_their_ending():
    assert sentences(" Kitchen light on.  Heating off! Is it?\nYes") == [
        "Kitchen light on.",
        "Heating off!",
        "Is it?",
        "Yes",
    ]


def test_an_utterance_is_one_session_committed_when_listening_ends(monkeypatch):
    heard: list[bytes] = []

    async def session(api_base, audio, *, rate, idle_s):
        assert (api_base, rate) == (f"{URL}/v1", 16000)
        while (chunk := await audio.get()) is not None:
            heard.append(chunk)
        yield False, "turn on"
        yield False, " the"
        yield True, " Turn on the light."

    monkeypatch.setattr(client, "transcribe_stream", session)

    async def scenario() -> None:
        node = SttNode(
            requires=[
                MessageSpec(name="mic", port="audio", dtype=DType.AUDIO),
                MessageSpec(name="listening", dtype=DType.BOOL),
                MessageSpec(name="url", dtype=DType.STR),
            ],
            provides=[
                MessageSpec(name="partial", dtype=DType.STR),
                MessageSpec(name="text", dtype=DType.STR),
            ],
        )
        node.assign_flow("voice", "stt")
        pipeline = Pipeline([node])
        store = Store()
        chunks = [
            store.put([bytes([n]) * 3200], "voice.mic", "audio/L16;rate=16000")
            for n in range(3)
        ]
        # Left over from before the engine started: not part of what is said.
        pipeline.state["voice.mic"] = chunks[0]
        await node.start()
        node._artifacts = store

        async def hear(audio: dict, listening: bool, url: str = URL) -> None:
            inputs = {"voice.mic": audio, "voice.listening": listening}
            await asyncio.to_thread(node.execute, inputs | {"voice.url": url})

        await hear(chunks[0], True)
        await hear(chunks[1], True)
        await hear(chunks[2], True)
        await hear(chunks[2], False)
        for _ in range(200):
            if "voice.text" in pipeline.state:
                break
            await asyncio.sleep(0.01)

        assert heard == [bytes([1]) * 3200, bytes([2]) * 3200]
        assert pipeline.state["voice.text"] == "Turn on the light."
        assert pipeline.state["voice.partial"] == "Turn on the light."
        assert node._session is None

        # No server, no session: the node says so instead.
        await hear(chunks[1], True, url="")
        assert node._session is None
        assert node._health == "degraded"

    asyncio.run(scenario())


def _wav(seconds: float) -> bytes:
    out = io.BytesIO()
    with wave.open(out, "wb") as clip:
        clip.setnchannels(1)
        clip.setsampwidth(2)
        clip.setframerate(24000)
        clip.writeframes(b"\0\0" * int(24000 * seconds))
    return out.getvalue()


def test_a_sentence_is_spoken_once_and_the_rest_of_a_reply_queues(monkeypatch):
    spoken: list[str] = []
    published: list[dict] = []

    async def speak(api_base, text, voice):
        spoken.append(text)
        return _wav(len(spoken))

    monkeypatch.setattr(client, "speak", speak)
    monkeypatch.setattr(TtsNode, "_say", lambda self, speech: published.append(speech))
    node = TtsNode(
        requires=[
            MessageSpec(name="reply", port="text", dtype=DType.STR),
            MessageSpec(name="url", dtype=DType.STR),
        ],
        provides=[MessageSpec(name="speech", dtype=DType.AUDIO)],
    )
    node.assign_flow("voice", "tts")
    Pipeline([node])
    node._artifacts = Store()

    node.execute({"voice.reply": "Kitchen light on. Heating off.", "voice.url": URL})
    # Long after both stopped playing.
    node._until = -10.0
    node.execute({"voice.reply": "Kitchen light on.", "voice.url": URL})

    assert spoken == ["Kitchen light on.", "Heating off."]
    assert ["append" in speech for speech in published] == [False, True, False]
    assert published[2] == published[0]
    assert published[1]["meta"] == {"seconds": 2.0}