"""Tests for executor group HA / drain controller (F10)."""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock, patch


from qremotex.executor_group import DrainController


# ──────────────────────────────────────────────── DrainController ──


class TestDrainController:
    def test_not_draining_by_default(self) -> None:
        ctrl = DrainController(active_jobs_fn=lambda: 0, on_drained=AsyncMock())
        assert ctrl.is_draining is False

    async def test_handle_drain_sets_draining(self) -> None:
        ctrl = DrainController(active_jobs_fn=lambda: 0, on_drained=AsyncMock())
        ctrl.handle_drain_message({"type": "drain", "reason": "rolling update"})
        assert ctrl.is_draining is True

    async def test_duplicate_drain_ignored(self) -> None:
        ctrl = DrainController(active_jobs_fn=lambda: 0, on_drained=AsyncMock())
        ctrl.handle_drain_message({"type": "drain"})
        # Second call should not raise or create another task
        ctrl.handle_drain_message({"type": "drain"})
        assert ctrl.is_draining is True

    async def test_immediate_drain_when_no_jobs(self) -> None:
        on_drained = AsyncMock()
        ctrl = DrainController(
            active_jobs_fn=lambda: 0,
            on_drained=on_drained,
            poll_interval_sec=0.01,
        )
        ctrl.handle_drain_message({"type": "drain"})
        # Give the drain task time to complete
        await asyncio.sleep(0.1)
        on_drained.assert_awaited_once()

    async def test_drain_waits_for_active_jobs(self) -> None:
        active = [3]

        def get_active() -> int:
            return active[0]

        on_drained = AsyncMock()
        ctrl = DrainController(
            active_jobs_fn=get_active,
            on_drained=on_drained,
            poll_interval_sec=0.01,
        )
        ctrl.handle_drain_message({"type": "drain"})

        # Still active — should not have called on_drained yet
        await asyncio.sleep(0.05)
        on_drained.assert_not_awaited()

        # Simulate jobs completing
        active[0] = 0
        await asyncio.sleep(0.05)
        on_drained.assert_awaited_once()

    async def test_drain_with_gradual_completion(self) -> None:
        active = [5]
        on_drained = AsyncMock()
        ctrl = DrainController(
            active_jobs_fn=lambda: active[0],
            on_drained=on_drained,
            poll_interval_sec=0.01,
        )
        ctrl.handle_drain_message({"type": "drain"})

        # Reduce jobs gradually
        await asyncio.sleep(0.03)
        active[0] = 2
        await asyncio.sleep(0.03)
        on_drained.assert_not_awaited()

        active[0] = 0
        await asyncio.sleep(0.05)
        on_drained.assert_awaited_once()

    async def test_on_drained_exception_handled(self) -> None:
        on_drained = AsyncMock(side_effect=RuntimeError("boom"))
        ctrl = DrainController(
            active_jobs_fn=lambda: 0,
            on_drained=on_drained,
            poll_interval_sec=0.01,
        )
        # Should not raise
        ctrl.handle_drain_message({"type": "drain"})
        await asyncio.sleep(0.1)
        on_drained.assert_awaited_once()


# ──────────────────────────────────────────────── Config ──


class TestExecutorGroupConfig:
    def test_default_no_group(self) -> None:
        with patch.dict("os.environ", {}, clear=False):
            import os
            os.environ.pop("QRAPTOR_EXECUTOR_GROUP", None)
            os.environ.setdefault("QRAPTOR_GATEWAY_URL", "wss://gw.example.com/ws/executor")
            os.environ.setdefault("QRAPTOR_API_KEY", "qrx_test_key123")
            from qremotex.config import ExecutorConfig
            config = ExecutorConfig()
            assert config.group == ""

    def test_group_from_env(self) -> None:
        with patch.dict("os.environ", {
            "QRAPTOR_EXECUTOR_GROUP": "us-east-gpu",
            "QRAPTOR_GATEWAY_URL": "wss://gw.example.com/ws/executor",
            "QRAPTOR_API_KEY": "qrx_test_key123",
        }):
            from qremotex.config import ExecutorConfig
            config = ExecutorConfig()
            assert config.group == "us-east-gpu"


# ──────────────────────────────────────────────── Hello message ──


class TestHelloMessageGroup:
    def test_hello_without_group(self) -> None:
        from qremotex.protocol.messages import hello_message
        msg = hello_message(
            executor_id="exec-1",
            api_key_id="ak_123",
            version="1.0.0",
            capabilities=["python:execute:unrestricted"],
            labels={},
            system_info={"os": "linux"},
        )
        assert "group" not in msg

    def test_hello_with_group(self) -> None:
        from qremotex.protocol.messages import hello_message
        msg = hello_message(
            executor_id="exec-1",
            api_key_id="ak_123",
            version="1.0.0",
            capabilities=["python:execute:unrestricted"],
            labels={},
            system_info={"os": "linux"},
            group="us-east-gpu",
        )
        assert msg["group"] == "us-east-gpu"

    def test_hello_empty_group_omitted(self) -> None:
        from qremotex.protocol.messages import hello_message
        msg = hello_message(
            executor_id="exec-1",
            api_key_id="ak_123",
            version="1.0.0",
            capabilities=["python:execute:unrestricted"],
            labels={},
            system_info={"os": "linux"},
            group="",
        )
        assert "group" not in msg


# ──────────────────────────────────────────────── DRAIN MessageType ──


class TestDrainMessageType:
    def test_drain_type_exists(self) -> None:
        from qremotex.protocol.messages import MessageType
        assert MessageType.DRAIN.value == "drain"
