from concurrent.futures import Future
from collections.abc import Sequence
from dataclasses import dataclass
import unittest

from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, PROTOCOL_VERSION, parse_request
from pi_dictate_service.recorder import Capture, RecorderError
from pi_dictate_service.server import DictateServiceCore, Outbound


class Clock:
    def __init__(self) -> None:
        self.now = 0.0

    def __call__(self) -> float:
        return self.now

    def advance(self, seconds: float) -> None:
        self.now += seconds


@dataclass
class PendingWork:
    future: Future[object]
    function: object
    args: tuple[object, ...]


class ControlledExecutor:
    def __init__(self) -> None:
        self.pending: list[PendingWork] = []

    def submit(self, function: object, *args: object) -> Future[object]:
        future: Future[object] = Future()
        self.pending.append(PendingWork(future, function, args))
        return future

    def run_next(self) -> None:
        work = self.pending.pop(0)
        try:
            result = work.function(*work.args)  # type: ignore[operator]
        except BaseException as error:
            work.future.set_exception(error)
        else:
            work.future.set_result(result)

    def shutdown(self, wait: bool = True, cancel_futures: bool = False) -> None:
        return None


class FakeModel:
    def __init__(self) -> None:
        self.loaded = False
        self.transcript = "dictated text"
        self.load_error: BaseException | None = None
        self.transcribe_error: BaseException | None = None
        self.samples: Sequence[float] | None = None

    def load(self) -> None:
        if self.load_error:
            raise self.load_error
        self.loaded = True

    def transcribe(self, samples: Sequence[float]) -> str:
        self.samples = samples
        if self.transcribe_error:
            raise self.transcribe_error
        return self.transcript


class FakeRecorder:
    def __init__(self) -> None:
        self.started = False
        self.cancelled = False
        self.capture = Capture(b"\x00\x00\x00\x40", -15, 15, b"")
        self.audio_level = 0.0
        self.stop_error: RecorderError | None = None
        self.unexpected: RecorderError | None = None

    def start(self) -> None:
        self.started = True

    def stop(self) -> Capture:
        if self.stop_error:
            raise self.stop_error
        return self.capture

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

    def poll_unexpected(self) -> RecorderError | None:
        failure, self.unexpected = self.unexpected, None
        return failure


class RecorderFactory:
    def __init__(self) -> None:
        self.instances: list[FakeRecorder] = []

    def __call__(self) -> FakeRecorder:
        recorder = FakeRecorder()
        self.instances.append(recorder)
        return recorder


class CoreTestCase(unittest.TestCase):
    def setUp(self) -> None:
        self.clock = Clock()
        self.executor = ControlledExecutor()
        self.model = FakeModel()
        self.recorders = RecorderFactory()
        self.core = DictateServiceCore(
            model=self.model,
            recorder_factory=self.recorders,
            executor=self.executor,
            clock=self.clock,
            idle_seconds=900,
        )

    def request(
        self,
        client: str,
        request_type: str,
        request_id: str,
        generation_id: str | None = None,
    ) -> list[dict[str, object]]:
        frame: dict[str, object] = {
            "version": PROTOCOL_VERSION,
            "type": request_type,
            "requestId": request_id,
        }
        if generation_id is not None:
            frame["generationId"] = generation_id
        self.core.handle_request(client, parse_request(frame))
        return self.frames(self.core.drain_events(), client)

    @staticmethod
    def frames(events: list[Outbound], client: str | None = None) -> list[dict[str, object]]:
        return [event.frame for event in events if client is None or event.client_id == client]

    def hello(self, client: str = "a") -> None:
        events = self.request(client, "hello", f"hello-{client}")
        self.assertEqual(events[0]["type"], "pong")

    def load_to_listening(self, client: str = "a", generation: str = "g1") -> FakeRecorder:
        self.hello(client)
        self.request(client, "start", "start", generation)
        self.executor.run_next()
        self.core.poll()
        events = self.frames(self.core.drain_events(), client)
        self.assertEqual(events[-1]["state"], "listening")
        self.assertEqual(events[-1]["requestId"], "start")
        return self.recorders.instances[-1]


class HandshakeAndOwnershipTests(CoreTestCase):
    def test_hello_is_required_and_valid_errors_keep_session_open(self) -> None:
        events = self.request("a", "start", "s1", "g1")
        self.assertEqual(events[0]["code"], "hello_required")

        self.hello("a")
        self.request("a", "start", "s2", "g1")
        self.assertEqual(len(self.executor.pending), 1)

    def test_one_generation_is_globally_owned(self) -> None:
        self.load_to_listening()
        self.hello("b")

        busy = self.request("b", "start", "s2", "g2")
        self.assertEqual(busy[0]["code"], "busy")
        not_owner = self.request("b", "stop", "x", "g1")
        self.assertEqual(not_owner[0]["code"], "not_owner")
        stale = self.request("a", "stop", "x2", "old")
        self.assertEqual(stale[0]["code"], "generation_mismatch")

        stopping = self.request("a", "stop", "x3", "g1")
        self.assertEqual(stopping[0]["state"], "stopping")


class CancellationTests(CoreTestCase):
    def test_cancel_during_load_is_immediate_and_blocks_until_worker_finishes(self) -> None:
        self.hello("a")
        self.request("a", "start", "start", "g1")

        cancelled = self.request("a", "cancel", "cancel", "g1")
        self.assertEqual(cancelled[0]["type"], "cancelled")
        self.hello("b")
        self.assertEqual(self.request("b", "start", "start-b", "g2")[0]["code"], "busy")

        self.executor.run_next()
        self.core.poll()
        self.assertEqual(self.recorders.instances, [])
        self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])

    def test_cancel_during_inference_suppresses_later_result(self) -> None:
        self.load_to_listening()
        self.request("a", "stop", "stop", "g1")
        self.executor.run_next()
        self.core.poll()
        transcribing = self.frames(self.core.drain_events(), "a")
        self.assertEqual(transcribing[0]["state"], "transcribing")

        cancelled = self.request("a", "cancel", "cancel", "g1")
        self.assertEqual(cancelled[0]["type"], "cancelled")
        self.executor.run_next()
        self.core.poll()
        self.assertNotIn("result", [event["type"] for event in self.frames(self.core.drain_events())])

    def test_disconnect_cancels_capture_and_tombstones_workers(self) -> None:
        recorder = self.load_to_listening()

        self.core.disconnect("a")
        self.assertEqual(len(self.executor.pending), 1)
        self.executor.run_next()
        self.core.poll()

        self.assertTrue(recorder.cancelled)
        self.assertEqual(self.core.active_generation_id, None)


class AudioLevelTests(CoreTestCase):
    def test_emits_owned_normalized_audio_levels_at_no_more_than_ten_hertz(self) -> None:
        recorder = self.load_to_listening()
        recorder.audio_level = 0.625

        events: list[Outbound] = []
        for _ in range(40):
            self.clock.advance(0.025)
            self.core.poll()
            events.extend(self.core.drain_events())

        self.assertGreaterEqual(len(events), 9)
        self.assertLessEqual(len(events), 10)
        self.assertEqual({event.client_id for event in events}, {"a"})
        self.assertEqual({event.frame["type"] for event in events}, {"level"})
        self.assertEqual({event.frame["level"] for event in events}, {0.625})

    def test_clamps_invalid_levels_and_stops_emitting_after_listening(self) -> None:
        recorder = self.load_to_listening()
        for raw_level, expected in ((2.0, 1.0), (-1.0, 0.0), (float("nan"), 0.0)):
            recorder.audio_level = raw_level
            self.clock.advance(0.1)
            self.core.poll()
            level = self.frames(self.core.drain_events(), "a")
            self.assertEqual(level[0]["level"], expected)

        self.request("a", "stop", "stop", "g1")
        self.clock.advance(0.2)
        self.core.poll()
        self.assertEqual(self.core.drain_events(), [])


class StopAndInferenceTests(CoreTestCase):
    def test_stop_drains_then_converts_before_inference_and_emits_result(self) -> None:
        self.load_to_listening()

        stopping = self.request("a", "stop", "stop", "g1")
        self.assertEqual(stopping[0]["state"], "stopping")
        self.assertEqual(stopping[0]["requestId"], "stop")
        self.executor.run_next()
        self.core.poll()
        transcribing = self.frames(self.core.drain_events(), "a")
        self.assertEqual(transcribing[0]["state"], "transcribing")

        self.executor.run_next()
        self.core.poll()
        completed = self.frames(self.core.drain_events(), "a")

        self.assertEqual(list(self.model.samples or []), [0.0, 0.5])
        self.assertEqual(completed[0]["type"], "result")
        self.assertEqual(completed[0]["transcript"], "dictated text")
        self.assertEqual(completed[1]["state"], "idle")

    def test_empty_transcript_is_a_result(self) -> None:
        self.model.transcript = ""
        self.load_to_listening()
        self.request("a", "stop", "stop", "g1")
        self.executor.run_next()
        self.core.poll()
        self.core.drain_events()
        self.executor.run_next()
        self.core.poll()

        completed = self.frames(self.core.drain_events(), "a")
        self.assertEqual(completed[0]["type"], "result")
        self.assertEqual(completed[0]["transcript"], "")

    def test_model_and_transcript_failures_are_bounded_terminal_errors(self) -> None:
        self.model.load_error = RuntimeError("CUDA unavailable")
        self.hello()
        self.request("a", "start", "start", "g1")
        self.executor.run_next()
        self.core.poll()
        failed = self.frames(self.core.drain_events(), "a")
        self.assertEqual(failed[0]["code"], "model_failure")
        self.assertEqual(failed[-1]["state"], "idle")

        model = FakeModel()
        model.loaded = True
        model.transcript = "x" * (MAX_TRANSCRIPT_BYTES + 1)
        core = DictateServiceCore(
            model=model,
            recorder_factory=self.recorders,
            executor=self.executor,
            clock=self.clock,
        )
        core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "hello", "requestId": "h"}))
        core.drain_events()
        core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "start", "requestId": "s", "generationId": "g"}))
        core.drain_events()
        core.handle_request("a", parse_request({"version": PROTOCOL_VERSION, "type": "stop", "requestId": "t", "generationId": "g"}))
        self.executor.run_next()
        core.poll()
        core.drain_events()
        self.executor.run_next()
        core.poll()
        failed = self.frames(core.drain_events(), "a")
        self.assertEqual(failed[0]["code"], "transcript_too_large")

    def test_unexpected_recorder_exit_is_terminal(self) -> None:
        recorder = self.load_to_listening()
        recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 3")

        self.core.poll()
        failed = self.frames(self.core.drain_events(), "a")

        self.assertEqual(failed[0]["code"], "recorder_exit")
        self.assertEqual(failed[-1]["state"], "idle")

    def test_stop_does_not_reclassify_an_already_exited_recorder_as_intentional(self) -> None:
        recorder = self.load_to_listening()
        recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 0")

        failed = self.request("a", "stop", "stop", "g1")

        self.assertEqual(failed[0]["code"], "recorder_exit")
        self.assertEqual(failed[-1]["state"], "idle")
        self.assertEqual(self.executor.pending, [])


class IdleTests(CoreTestCase):
    def test_only_valid_requests_reset_idle_and_workers_suspend_then_restart_timer(self) -> None:
        self.core.accepted("a")
        self.clock.advance(899)
        self.assertFalse(self.core.should_exit())
        self.request("a", "hello", "h")
        self.clock.advance(899)
        self.assertFalse(self.core.should_exit())

        self.request("a", "start", "s", "g")
        self.clock.advance(901)
        self.assertFalse(self.core.should_exit())
        self.request("a", "cancel", "c", "g")
        self.executor.run_next()
        self.core.poll()
        self.core.drain_events()
        self.assertFalse(self.core.should_exit())

        self.clock.advance(899)
        self.assertFalse(self.core.should_exit())
        self.clock.advance(1)
        self.assertTrue(self.core.should_exit())


if __name__ == "__main__":
    unittest.main()