from __future__ import annotations

import asyncio
import unittest
from types import SimpleNamespace

from models import AssistantConfig
from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    BotStoppedSpeakingFrame,
    TTSSpeakFrame,
)
import services.brains  # Initialize the brain registry before realtime imports.
from services.fixed_speech import FixedSpeechOutput
from services.message_stage import MESSAGE_CONFIRMATION_CONTEXT_MARKER
from services.pipecat.call_lifecycle import (
    CallEndCoordinator,
    FixedSpeechPlaybackMarkerFrame,
)
from services.pipecat.transports import build_ws_transport
from services.runtime_variables import DynamicVariableStore
from services.workflow.realtime import WorkflowRealtimeController
from services.workflow_engine import WorkflowEngine


class FixedSpeechPlaybackTest(unittest.IsolatedAsyncioTestCase):
    async def test_pipeline_speech_queues_marker_immediately_after_tts(self):
        queued = []

        async def queue_end(_reason: str) -> None:
            pass

        async def queue_frame(frame) -> None:
            queued.append(frame)

        call_end = CallEndCoordinator(queue_end)
        output = FixedSpeechOutput(
            DynamicVariableStore({}),
            SimpleNamespace(call_end=call_end, queue_frame=queue_frame),
        )

        completion = await output.speak(
            "固定结束语",
            source="test",
            record_history=False,
        )

        self.assertIsInstance(queued[0], TTSSpeakFrame)
        self.assertIsInstance(queued[1], FixedSpeechPlaybackMarkerFrame)
        self.assertIs(queued[1].completion, completion)

    async def test_websocket_output_resolves_marker(self):
        async def queue_end(_reason: str) -> None:
            pass

        call_end = CallEndCoordinator(queue_end)
        completion = call_end.track_speech()
        marker = FixedSpeechPlaybackMarkerFrame(completion=completion)
        websocket = SimpleNamespace(headers={})
        output = build_ws_transport(websocket).output()

        await output.write_transport_frame(marker)

        self.assertTrue(completion.done())

    async def test_realtime_end_ignores_unrelated_speech_boundaries(self):
        graph = {
            "specVersion": 3,
            "settings": {},
            "nodes": [
                {"id": "start", "type": "start", "data": {}},
                {
                    "id": "end",
                    "type": "end",
                    "data": {"message": "感谢来电，再见。", "scope": "session"},
                },
            ],
            "edges": [],
        }
        reasons = []
        queued = []

        async def queue_end(reason: str) -> None:
            reasons.append(reason)

        async def queue_frame(frame) -> None:
            queued.append(frame)

        class FakeRealtime:
            def __init__(self):
                self.provider_completion = None

            async def update_session(self, _instructions, _tools):
                pass

            async def speak_fixed(self, _text, *, suppress_transcript=True):
                self.provider_completion = (
                    asyncio.get_running_loop().create_future()
                )
                return self.provider_completion

        call_end = CallEndCoordinator(queue_end)
        realtime = FakeRealtime()
        controller = WorkflowRealtimeController(
            cfg=AssistantConfig(type="workflow", graph=graph),
            engine=WorkflowEngine(graph),
            store=DynamicVariableStore({}),
            runtime=SimpleNamespace(
                realtime=realtime,
                queue_frame=queue_frame,
                call_end=call_end,
                session_id="test-session",
                client_tools=None,
                set_input_enabled=lambda _enabled: None,
                capture_image=None,
            ),
        )

        end_task = asyncio.create_task(controller._enter_end("end"))
        while realtime.provider_completion is None:
            await asyncio.sleep(0)

        await call_end.observe(BotStartedSpeakingFrame())
        await call_end.observe(BotStoppedSpeakingFrame())
        self.assertEqual(reasons, [])
        self.assertFalse(end_task.done())

        realtime.provider_completion.set_result(None)
        marker = None
        while marker is None:
            await asyncio.sleep(0)
            marker = next(
                (
                    frame
                    for frame in queued
                    if isinstance(frame, FixedSpeechPlaybackMarkerFrame)
                ),
                None,
            )
        await marker.completion.mark_played()
        await end_task

        self.assertEqual(reasons, ["workflow_completed"])

    async def test_realtime_confirmation_is_appended_without_running_model(self):
        graph = {
            "specVersion": 3,
            "settings": {},
            "nodes": [
                {
                    "id": "message",
                    "type": "message",
                    "data": {
                        "title": "办理提示",
                        "message": "确认后开始办理。",
                        "confirmLabel": "继续办理",
                        "completionPolicy": "confirmation",
                    },
                }
            ],
            "edges": [],
        }
        appended = []
        input_states = []

        class FakeRealtime:
            async def send_text(self, text, *, run_immediately=True):
                appended.append((text, run_immediately))

        class FakeClientTools:
            async def call(self, *_args, **_kwargs):
                return {"status": "ok", "data": {"action": "confirmed"}}

        async def queue_frame(_frame):
            pass

        controller = WorkflowRealtimeController(
            cfg=AssistantConfig(type="workflow", graph=graph),
            engine=WorkflowEngine(graph),
            store=DynamicVariableStore({}),
            runtime=SimpleNamespace(
                realtime=FakeRealtime(),
                queue_frame=queue_frame,
                call_end=SimpleNamespace(ending=False),
                session_id="test-session",
                client_tools=FakeClientTools(),
                set_input_enabled=input_states.append,
                capture_image=None,
            ),
        )

        succeeded = await controller._enter_message("message")

        self.assertTrue(succeeded)
        self.assertEqual(input_states, [False, True])
        self.assertEqual(len(appended), 1)
        self.assertFalse(appended[0][1])
        self.assertIn(MESSAGE_CONFIRMATION_CONTEXT_MARKER, appended[0][0])
        self.assertIn("点击“继续办理”", appended[0][0])


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