import asyncio
import signal
import unittest
from unittest.mock import AsyncMock, patch

from rex.client import action_manager, receive_actions, run_command
from rex.client.action_manager import ClientAction


class FakeWebSocket:
    def __init__(self, messages: list[str]) -> None:
        self.messages = iter(messages)

    def __aiter__(self):
        return self

    async def __anext__(self) -> str:
        await asyncio.sleep(0)
        try:
            return next(self.messages)
        except StopIteration:
            raise StopAsyncIteration


class ClientTests(unittest.IsolatedAsyncioTestCase):
    async def test_run_command_waits_for_process(self) -> None:
        process = AsyncMock()
        process.wait.return_value = 0

        with patch("rex.client.asyncio.create_subprocess_shell", return_value=process):
            await run_command("test command")

        process.wait.assert_awaited_once()

    async def test_cancelling_command_terminates_process_group(self) -> None:
        process = AsyncMock()
        process.pid = 123
        process.returncode = None
        wait_started = asyncio.Event()
        wait_calls = 0

        async def wait() -> int:
            nonlocal wait_calls
            wait_calls += 1
            if wait_calls == 1:
                wait_started.set()
                await asyncio.Event().wait()
            return 0

        process.wait.side_effect = wait

        with (
            patch("rex.client.asyncio.create_subprocess_shell", return_value=process),
            patch("rex.client.os.name", "posix"),
            patch("rex.client.os.killpg") as killpg,
        ):
            task = asyncio.create_task(run_command("test command"))
            await wait_started.wait()
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)

        killpg.assert_called_once_with(123, signal.SIGTERM)
        self.assertEqual(process.wait.await_count, 2)

    async def test_disconnect_cancels_running_commands(self) -> None:
        action = ClientAction(name="wake", command="test command")
        command_started = asyncio.Event()
        command_cancelled = asyncio.Event()

        async def running_command(_: str) -> None:
            command_started.set()
            try:
                await asyncio.Event().wait()
            finally:
                command_cancelled.set()

        websocket = FakeWebSocket(['{"name":"wake"}'])
        with (
            patch.object(action_manager, "get_action", return_value=action),
            patch("rex.client.run_command", side_effect=running_command),
        ):
            await receive_actions(websocket)  # type: ignore[arg-type]

        self.assertTrue(command_started.is_set())
        self.assertTrue(command_cancelled.is_set())


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