import asyncio
import os
import sys
import unittest
from pathlib import Path
from unittest.mock import patch


BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from rpa_engine.douyin_im.traffic_control import (
    GlobalSendQueue,
    TrafficController,
    get_traffic_controller,
)


class GlobalSendQueueTests(unittest.IsolatedAsyncioTestCase):
    async def _make_queue(self, interval_seconds: float = 0.0) -> GlobalSendQueue:
        queue = GlobalSendQueue(interval_seconds=interval_seconds)
        self.addAsyncCleanup(queue.stop)
        return queue

    async def _wait_for_queued_count(
        self,
        queue: GlobalSendQueue,
        expected: int,
        timeout: float = 0.5,
    ) -> None:
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout
        while loop.time() < deadline:
            if (await queue.snapshot())["queued_count"] == expected:
                return
            await asyncio.sleep(0.001)
        self.assertEqual((await queue.snapshot())["queued_count"], expected)

    async def test_global_send_concurrency_is_one_across_accounts(self):
        queue = await self._make_queue()
        active = 0
        maximum_active = 0

        async def operation(index: int) -> int:
            nonlocal active, maximum_active
            active += 1
            maximum_active = max(maximum_active, active)
            await asyncio.sleep(0.015)
            active -= 1
            return index

        tasks = [
            asyncio.create_task(
                queue.submit(
                    account_id=(index % 3) + 1,
                    operation=lambda index=index: operation(index),
                    description=f"send-{index}",
                )
            )
            for index in range(9)
        ]

        results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=1.0)

        self.assertEqual(sorted(results), list(range(9)))
        self.assertEqual(maximum_active, 1)
        self.assertEqual(active, 0)

    async def test_same_account_keeps_fifo_order(self):
        queue = await self._make_queue()
        calls: list[int] = []

        async def operation(index: int) -> int:
            calls.append(index)
            await asyncio.sleep(0)
            return index

        tasks = [
            asyncio.create_task(
                queue.submit(
                    account_id=101,
                    operation=lambda index=index: operation(index),
                    description=str(index),
                )
            )
            for index in range(6)
        ]

        results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.5)

        self.assertEqual(calls, list(range(6)))
        self.assertEqual(results, list(range(6)))

    async def test_busy_accounts_are_dispatched_round_robin(self):
        queue = await self._make_queue()
        calls: list[str] = []
        first_started = asyncio.Event()
        release_first = asyncio.Event()

        async def first_operation() -> str:
            calls.append("a1")
            first_started.set()
            await release_first.wait()
            return "a1"

        async def operation(name: str) -> str:
            calls.append(name)
            await asyncio.sleep(0)
            return name

        first = asyncio.create_task(queue.submit(1, first_operation, "a1"))
        await asyncio.wait_for(first_started.wait(), timeout=0.2)

        # Account 1 is deliberately noisy.  Accounts 2 and 3 must each get a
        # turn before account 1 is allowed to dispatch its next waiting job.
        remaining_specs = [
            (1, "a2"),
            (1, "a3"),
            (2, "b1"),
            (2, "b2"),
            (3, "c1"),
        ]
        remaining = [
            asyncio.create_task(
                queue.submit(
                    account_id,
                    lambda name=name: operation(name),
                    description=name,
                )
            )
            for account_id, name in remaining_specs
        ]
        await self._wait_for_queued_count(queue, expected=len(remaining))

        release_first.set()
        results = await asyncio.wait_for(
            asyncio.gather(first, *remaining),
            timeout=0.5,
        )

        self.assertEqual(calls, ["a1", "b1", "c1", "a2", "b2", "a3"])
        self.assertCountEqual(results, ["a1", "a2", "a3", "b1", "b2", "c1"])

    async def test_send_starts_are_separated_by_configured_interval(self):
        interval = 0.08
        queue = await self._make_queue(interval_seconds=interval)
        loop = asyncio.get_running_loop()
        started_at: list[float] = []

        async def operation() -> None:
            started_at.append(loop.time())

        tasks = [
            asyncio.create_task(queue.submit(index + 1, operation, str(index)))
            for index in range(3)
        ]
        await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.7)

        self.assertEqual(len(started_at), 3)
        gaps = [later - earlier for earlier, later in zip(started_at, started_at[1:])]
        for gap in gaps:
            # The Windows selector clock may wake one ~15.6 ms tick early.
            self.assertGreaterEqual(gap, interval - 0.025)

    async def test_operation_failure_does_not_block_later_job(self):
        queue = await self._make_queue()
        calls: list[str] = []

        async def failing_operation() -> None:
            calls.append("failed")
            raise ValueError("expected failure")

        async def following_operation() -> str:
            calls.append("continued")
            return "ok"

        failed = asyncio.create_task(queue.submit(1, failing_operation, "failed"))
        continued = asyncio.create_task(queue.submit(2, following_operation, "continued"))

        results = await asyncio.wait_for(
            asyncio.gather(failed, continued, return_exceptions=True),
            timeout=0.2,
        )

        self.assertIsInstance(results[0], ValueError)
        self.assertEqual(str(results[0]), "expected failure")
        self.assertEqual(results[1], "ok")
        self.assertEqual(calls, ["failed", "continued"])

    async def test_cancelling_waiting_submit_prevents_operation(self):
        queue = await self._make_queue()
        first_started = asyncio.Event()
        release_first = asyncio.Event()
        cancelled_operation_called = asyncio.Event()

        async def blocking_operation() -> None:
            first_started.set()
            await release_first.wait()

        async def must_not_run() -> None:
            cancelled_operation_called.set()

        blocker = asyncio.create_task(queue.submit(1, blocking_operation, "blocker"))
        await asyncio.wait_for(first_started.wait(), timeout=0.2)
        waiting = asyncio.create_task(queue.submit(2, must_not_run, "cancelled"))
        await self._wait_for_queued_count(queue, expected=1)

        waiting.cancel()
        with self.assertRaises(asyncio.CancelledError):
            await waiting
        release_first.set()
        await asyncio.wait_for(blocker, timeout=0.2)
        await asyncio.sleep(0.02)

        self.assertFalse(cancelled_operation_called.is_set())
        self.assertEqual((await queue.snapshot())["pending_count"], 0)

    async def test_cancelling_active_submit_waits_for_real_result_without_overlap(self):
        queue = await self._make_queue()
        active = 0
        maximum_active = 0
        first_started = asyncio.Event()
        release_first = asyncio.Event()
        second_started = asyncio.Event()

        async def first_operation() -> str:
            nonlocal active, maximum_active
            active += 1
            maximum_active = max(maximum_active, active)
            first_started.set()
            await release_first.wait()
            active -= 1
            return "delivered"

        async def second_operation() -> str:
            nonlocal active, maximum_active
            active += 1
            maximum_active = max(maximum_active, active)
            second_started.set()
            await asyncio.sleep(0)
            active -= 1
            return "next"

        first = asyncio.create_task(queue.submit(1, first_operation, "active"))
        await asyncio.wait_for(first_started.wait(), timeout=0.2)
        second = asyncio.create_task(queue.submit(2, second_operation, "next"))
        await self._wait_for_queued_count(queue, expected=1)

        first.cancel()
        await asyncio.sleep(0.02)

        # Once the network operation has begun, cancelling its caller must not
        # release the single lane or report a false cancellation to the caller.
        self.assertFalse(first.done())
        self.assertFalse(second_started.is_set())
        self.assertEqual(active, 1)

        release_first.set()
        self.assertEqual(await asyncio.wait_for(first, timeout=0.2), "delivered")
        self.assertEqual(await asyncio.wait_for(second, timeout=0.2), "next")
        self.assertTrue(second_started.is_set())
        self.assertEqual(maximum_active, 1)
        self.assertEqual(active, 0)

    async def test_cancel_account_removes_waiting_jobs_without_running_them(self):
        queue = await self._make_queue()
        blocker_started = asyncio.Event()
        release_blocker = asyncio.Event()
        removed_calls: list[str] = []

        async def blocking_operation() -> None:
            blocker_started.set()
            await release_blocker.wait()

        async def removed_operation(name: str) -> None:
            removed_calls.append(name)

        blocker = asyncio.create_task(queue.submit(1, blocking_operation, "blocker"))
        await asyncio.wait_for(blocker_started.wait(), timeout=0.2)
        removed = [
            asyncio.create_task(
                queue.submit(
                    77,
                    lambda name=name: removed_operation(name),
                    name,
                )
            )
            for name in ("waiting-1", "waiting-2")
        ]
        await self._wait_for_queued_count(queue, expected=2)

        self.assertEqual(await queue.cancel_account(77), 2)
        results = await asyncio.wait_for(
            asyncio.gather(*removed, return_exceptions=True),
            timeout=0.1,
        )
        snapshot = await queue.snapshot()

        self.assertTrue(all(isinstance(result, asyncio.CancelledError) for result in results))
        self.assertEqual(removed_calls, [])
        self.assertNotIn(77, snapshot["per_account"])
        self.assertEqual(snapshot["queued_count"], 0)
        self.assertEqual(snapshot["active_account_id"], 1)

        release_blocker.set()
        await asyncio.wait_for(blocker, timeout=0.2)
        await asyncio.sleep(0)
        self.assertEqual((await queue.snapshot())["pending_count"], 0)

    async def test_cancel_account_releases_active_caller_but_drains_lane(self):
        queue = await self._make_queue()
        active = 0
        maximum_active = 0
        active_started = asyncio.Event()
        release_active = asyncio.Event()
        active_drained = asyncio.Event()
        next_started = asyncio.Event()
        calls: list[str] = []

        async def active_operation() -> str:
            nonlocal active, maximum_active
            active += 1
            maximum_active = max(maximum_active, active)
            calls.append("active-start")
            active_started.set()
            await release_active.wait()
            calls.append("active-drained")
            active -= 1
            active_drained.set()
            return "discarded-result"

        async def next_operation() -> str:
            nonlocal active, maximum_active
            active += 1
            maximum_active = max(maximum_active, active)
            calls.append("next")
            next_started.set()
            await asyncio.sleep(0)
            active -= 1
            return "next-result"

        active_submit = asyncio.create_task(
            queue.submit(88, active_operation, "active-account")
        )
        await asyncio.wait_for(active_started.wait(), timeout=0.2)
        next_submit = asyncio.create_task(queue.submit(89, next_operation, "other-account"))
        await self._wait_for_queued_count(queue, expected=1)

        self.assertEqual(await queue.cancel_account(88), 1)
        cancelled_result = await asyncio.wait_for(
            asyncio.gather(active_submit, return_exceptions=True),
            timeout=0.1,
        )

        self.assertIsInstance(cancelled_result[0], asyncio.CancelledError)
        self.assertTrue(active_submit.cancelled())
        self.assertFalse(active_drained.is_set())
        self.assertFalse(next_started.is_set())
        self.assertEqual(active, 1)

        # The cancelled account's network operation still owns the lane until
        # it actually drains; only then can a different account begin.
        release_active.set()
        self.assertEqual(
            await asyncio.wait_for(next_submit, timeout=0.2),
            "next-result",
        )
        self.assertTrue(active_drained.is_set())
        self.assertEqual(calls, ["active-start", "active-drained", "next"])
        self.assertEqual(maximum_active, 1)
        self.assertEqual(active, 0)

        # A fresh submission proves the dispatcher remains usable after the
        # cancelled active job and its queued successor have both completed.
        async def recovered_operation() -> str:
            calls.append("recovered")
            return "recovered-result"

        self.assertEqual(
            await asyncio.wait_for(
                queue.submit(90, recovered_operation, "recovered"),
                timeout=0.2,
            ),
            "recovered-result",
        )
        self.assertEqual(calls[-1], "recovered")
        self.assertEqual((await queue.snapshot())["pending_count"], 0)


class BackgroundTrafficLimitTests(unittest.IsolatedAsyncioTestCase):
    async def _wait_until(self, predicate) -> None:
        while not predicate():
            await asyncio.sleep(0.001)

    async def test_background_slot_respects_configured_concurrency_limit(self):
        with patch.dict(
            os.environ,
            {"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "2"},
        ):
            controller = TrafficController()
        self.addAsyncCleanup(controller.stop)

        active = 0
        maximum_active = 0
        entered = 0
        all_finished = asyncio.Event()

        async def request(account_id: int) -> None:
            nonlocal active, maximum_active, entered
            async with controller.background_slot(account_id, "poll"):
                active += 1
                entered += 1
                maximum_active = max(maximum_active, active)
                await asyncio.sleep(0.02)
                active -= 1
                if entered == 8 and active == 0:
                    all_finished.set()

        tasks = [asyncio.create_task(request(index + 1)) for index in range(8)]
        await asyncio.wait_for(asyncio.gather(*tasks), timeout=0.5)

        self.assertEqual(entered, 8)
        self.assertEqual(maximum_active, 2)
        self.assertEqual(active, 0)
        self.assertEqual(controller.background_active, 0)
        self.assertEqual(controller.background_waiting, 0)

    async def test_nested_background_slot_in_same_task_is_reentrant(self):
        with patch.dict(
            os.environ,
            {"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
        ):
            controller = TrafficController()
        self.addAsyncCleanup(controller.stop)

        tasks_seen: list[asyncio.Task] = []

        async def nested_request() -> None:
            tasks_seen.append(asyncio.current_task())
            async with controller.background_slot(1, "outer"):
                self.assertEqual(controller.background_active, 1)
                self.assertEqual(controller.background_waiting, 0)
                self.assertEqual(controller._background._value, 0)

                tasks_seen.append(asyncio.current_task())
                async with controller.background_slot(1, "inner"):
                    tasks_seen.append(asyncio.current_task())
                    self.assertEqual(controller.background_active, 1)
                    self.assertEqual(controller.background_waiting, 0)
                    self.assertEqual(controller._background._value, 0)

                self.assertEqual(controller.background_active, 1)
                self.assertEqual(controller._background._value, 0)

            self.assertEqual(controller.background_active, 0)
            self.assertEqual(controller.background_waiting, 0)
            self.assertEqual(controller._background._value, 1)

        await asyncio.wait_for(nested_request(), timeout=0.2)

        self.assertEqual(len(tasks_seen), 3)
        self.assertTrue(all(task is tasks_seen[0] for task in tasks_seen))

    async def test_startup_traffic_leaves_one_slot_for_recurring_work(self):
        with patch.dict(
            os.environ,
            {"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "3"},
        ):
            controller = TrafficController()
        self.addAsyncCleanup(controller.stop)

        entered: list[str] = []
        release = asyncio.Event()

        async def request(name: str, *, startup: bool = False) -> None:
            async with controller.background_slot(1, name, startup=startup):
                entered.append(name)
                await release.wait()

        starts = [
            asyncio.create_task(request(f"startup-{index}", startup=True))
            for index in range(3)
        ]
        while len(entered) < 2:
            await asyncio.sleep(0)
        await asyncio.sleep(0.01)

        # A flood of startups may occupy at most capacity - 1 slots, so a
        # recurring poll still gets in while the fleet is coming online.
        self.assertEqual(len(entered), 2)
        self.assertEqual(controller.background_startup_active, 2)

        poll = asyncio.create_task(request("poll"))
        await asyncio.wait_for(
            self._wait_until(lambda: "poll" in entered),
            timeout=0.5,
        )

        release.set()
        await asyncio.wait_for(asyncio.gather(*starts, poll), timeout=0.5)
        self.assertEqual(controller.background_active, 0)
        self.assertEqual(controller.background_startup_active, 0)

    async def test_recurring_work_is_not_deferred_indefinitely_by_startups(self):
        """Pending startup work may delay a recurring poll, never block it.

        Starting several hundred accounts keeps startup requests queued for the
        whole run.  Yielding to that queue without a deadline left every hosted
        account silent until the last account had finished coming online.
        """
        with patch.dict(
            os.environ,
            {
                "KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1",
                "KEFU_BACKGROUND_NORMAL_MAX_DEFER_SECONDS": "0.02",
            },
        ):
            controller = TrafficController()
        self.addAsyncCleanup(controller.stop)

        entered: list[str] = []
        release = asyncio.Event()

        async def poll() -> None:
            async with controller.background_slot(1, "conversation poll"):
                entered.append("poll")
                await release.wait()

        # Stands in for a batch whose startup requests never stop arriving.
        controller._background_startup_clear.clear()

        task = asyncio.create_task(poll())
        await asyncio.wait_for(
            self._wait_until(lambda: entered == ["poll"]),
            timeout=1.0,
        )

        release.set()
        await asyncio.wait_for(task, timeout=0.5)

    async def test_startup_request_is_not_buried_behind_normal_backlog(self):
        with patch.dict(
            os.environ,
            {"KEFU_BACKGROUND_NETWORK_CONCURRENCY": "1"},
        ):
            controller = TrafficController()
        self.addAsyncCleanup(controller.stop)

        entered: list[str] = []
        releases = {
            name: asyncio.Event()
            for name in ("active", "normal-1", "normal-2", "startup")
        }

        async def request(name: str, *, startup: bool = False) -> None:
            async with controller.background_slot(
                1,
                name,
                startup=startup,
            ):
                entered.append(name)
                await releases[name].wait()

        active = asyncio.create_task(request("active"))
        while entered != ["active"]:
            await asyncio.sleep(0)

        normal_one = asyncio.create_task(request("normal-1"))
        normal_two = asyncio.create_task(request("normal-2"))
        # Let one normal request reach the shared semaphore while the other is
        # held at normal admission, then add the priority startup request.
        await asyncio.sleep(0)
        await asyncio.sleep(0)
        startup = asyncio.create_task(request("startup", startup=True))
        await asyncio.sleep(0)

        releases["active"].set()
        while len(entered) < 2:
            await asyncio.sleep(0)
        self.assertEqual(entered[:2], ["active", "normal-1"])

        releases["normal-1"].set()
        while len(entered) < 3:
            await asyncio.sleep(0)
        self.assertEqual(entered[:3], ["active", "normal-1", "startup"])

        releases["startup"].set()
        while len(entered) < 4:
            await asyncio.sleep(0)
        releases["normal-2"].set()
        await asyncio.wait_for(
            asyncio.gather(active, normal_one, normal_two, startup),
            timeout=0.2,
        )

        self.assertEqual(
            entered,
            ["active", "normal-1", "startup", "normal-2"],
        )
        self.assertEqual(controller.background_active, 0)
        self.assertEqual(controller.background_waiting, 0)
        self.assertEqual(controller.background_startup_active, 0)
        self.assertEqual(controller.background_startup_waiting, 0)


class TrafficControllerLoopIsolationTests(unittest.TestCase):
    def test_get_traffic_controller_does_not_reuse_asyncio_primitives(self):
        async def capture_controller_state():
            controller = get_traffic_controller()
            return {
                "loop": asyncio.get_running_loop(),
                "controller": controller,
                "send_lock": controller.send_queue._state_lock,
                "send_wake": controller.send_queue._wake,
                "background": controller._background,
                "browser": controller._browser,
            }

        # IsolatedAsyncioTestCase uses a fresh asyncio.Runner per test.  Keep
        # both closed loop objects alive so the WeakKeyDictionary is exercised
        # with two distinct keys rather than relying on garbage collection.
        with asyncio.Runner() as first_runner:
            first = first_runner.run(capture_controller_state())
        with asyncio.Runner() as second_runner:
            second = second_runner.run(capture_controller_state())

        for key in (
            "loop",
            "controller",
            "send_lock",
            "send_wake",
            "background",
            "browser",
        ):
            self.assertIsNot(first[key], second[key], key)


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