"""Unit tests for prefill bootstrap queue compaction."""

import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=6, suite="stage-a-test-cpu")


class FakeSender:
    def __init__(self, should_fail=False):
        self.should_fail = should_fail
        self.init_calls = []

    def init(self, num_pages, metadata_buffer_index):
        self.init_calls.append((num_pages, metadata_buffer_index))

    def failure_exception(self):
        if self.should_fail:
            raise RuntimeError("boom")


class TestPrefillBootstrapQueue(CustomTestCase):
    def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
        req = MagicMock()
        req.rid = rid
        req.bootstrap_room = bootstrap_room
        req.origin_input_ids = origin_input_ids
        req.disagg_kv_sender = sender
        req.return_logprob = False
        req.metadata_buffer_index = -1
        req.time_stats = SimpleNamespace(
            set_bootstrap_done_time=MagicMock(),
            set_wait_queue_entry_time=MagicMock(),
            trace_ctx=SimpleNamespace(abort=MagicMock()),
        )
        return req

    def _build_queue(self, allocator_size=4):
        queue = PrefillBootstrapQueue.__new__(PrefillBootstrapQueue)
        queue.req_to_metadata_buffer_idx_allocator = ReqToMetadataIdxAllocator(
            size=allocator_size
        )
        queue.token_to_kv_pool = SimpleNamespace(page_size=4)
        queue.scheduler = SimpleNamespace(
            stream_output=MagicMock(),
            enable_metrics=False,
            enable_hicache_storage=False,
            attn_cp_cpu_group=MagicMock(),
            attn_tp_cpu_group=MagicMock(),
        )
        queue.tp_rank = 0
        queue.queue = []
        return queue

    def test_pop_bootstrapped_compacts_queue_in_one_pass(self):
        queue = self._build_queue(allocator_size=2)

        waiting = self._make_req("wait", 1, list(range(8)), FakeSender())
        ready = self._make_req("ready", 2, list(range(10)), FakeSender())
        failed = self._make_req("failed", 3, list(range(6)), FakeSender(True))
        queue.queue = [waiting, ready, failed]

        with patch(
            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
            return_value=[
                KVPoll.Bootstrapping,
                KVPoll.WaitingForInput,
                KVPoll.Failed,
            ],
        ):
            bootstrapped, failed_reqs = queue.pop_bootstrapped(return_failed_reqs=True)

        self.assertEqual(bootstrapped, [ready])
        self.assertEqual(failed_reqs, [failed])
        self.assertEqual(queue.queue, [waiting])
        self.assertEqual(ready.metadata_buffer_index, 0)
        self.assertEqual(ready.disagg_kv_sender.init_calls, [(3, 0)])

    def test_pop_bootstrapped_stops_on_metadata_exhaustion(self):
        queue = self._build_queue(allocator_size=1)

        first_ready = self._make_req("ready-1", 1, list(range(8)), FakeSender())
        second_ready = self._make_req("ready-2", 2, list(range(12)), FakeSender())
        trailing = self._make_req("trailing", 3, list(range(4)), FakeSender())
        queue.queue = [first_ready, second_ready, trailing]

        with patch(
            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
            return_value=[
                KVPoll.WaitingForInput,
                KVPoll.WaitingForInput,
                KVPoll.Bootstrapping,
            ],
        ):
            bootstrapped = queue.pop_bootstrapped()

        self.assertEqual(bootstrapped, [first_ready])
        self.assertEqual(queue.queue, [second_ready, trailing])
        self.assertEqual(first_ready.disagg_kv_sender.init_calls, [(2, 0)])
        self.assertEqual(second_ready.disagg_kv_sender.init_calls, [])

    def test_pop_bootstrapped_preserves_unchecked_entries(self):
        queue = self._build_queue(allocator_size=2)

        skipped = self._make_req("skip", 1, list(range(8)), FakeSender())
        checked = self._make_req("check", 2, list(range(10)), FakeSender())
        queue.queue = [skipped, checked]

        with patch(
            "sglang.srt.disaggregation.prefill.poll_and_all_reduce_attn_cp_tp_group",
            return_value=[KVPoll.WaitingForInput, KVPoll.WaitingForInput],
        ):
            bootstrapped = queue.pop_bootstrapped(rids_to_check=["check"])

        self.assertEqual(bootstrapped, [checked])
        self.assertEqual(queue.queue, [skipped])
        self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
        self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])


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