from __future__ import annotations

import asyncio

import pytest

from yuxi.channels.services.context import ChatAbortEntry, ChatRunBuffer
from yuxi.channels.router import MessageRouter


class TestMessageRouterRuntime:
    def test_chat_abort_controllers_initialized(self):
        router = MessageRouter()
        assert router.chat_abort_controllers == {}
        assert router.chat_run_buffers == {}

    def test_abort_chat_returns_false_for_unknown_run(self):
        router = MessageRouter()
        assert router.abort_chat("nonexistent") is False

    @pytest.mark.asyncio
    async def test_abort_chat_cancels_running_task(self):
        router = MessageRouter()
        run_id = "test:msg-1"

        async def slow_task():
            await asyncio.sleep(10)

        loop = asyncio.get_running_loop()
        task = loop.create_task(slow_task())
        router.chat_abort_controllers[run_id] = ChatAbortEntry(task=task)

        assert router.abort_chat(run_id) is True
        await asyncio.sleep(0)
        assert task.cancelled()

    def test_chat_run_buffer_registration(self):
        router = MessageRouter()
        buffer = ChatRunBuffer(run_id="test:msg-2")
        router.chat_run_buffers["test:msg-2"] = buffer

        buffer.append_chunk("hello")
        assert router.chat_run_buffers["test:msg-2"].get_full_text() == "hello"

    def test_chat_abort_entry_cleanup_on_abort(self):
        router = MessageRouter()
        run_id = "test:msg-3"

        async def noop():
            pass

        loop = asyncio.new_event_loop()
        task = loop.create_task(noop())
        router.chat_abort_controllers[run_id] = ChatAbortEntry(task=task)

        router.abort_chat(run_id)

        router.chat_abort_controllers.pop(run_id, None)
        assert run_id not in router.chat_abort_controllers
        loop.close()