"""Tests for scheduler load accounting used by DP dispatch."""

from types import SimpleNamespace

# Import scheduler first to mirror production import order and avoid the
# scheduler <-> metrics circular import during isolated unit collection.
import sglang.srt.managers.scheduler  # noqa: F401

from sglang.srt.disaggregation.decode import DecodeRequest
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.observability.scheduler_metrics_mixin import SchedulerMetricsMixin
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

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


class FakeReq:
    def __init__(self, rid: str, prompt_tokens: int, output_tokens: int = 0):
        self.rid = rid
        self.origin_input_ids = list(range(prompt_tokens))
        self.origin_input_ids_unpadded = self.origin_input_ids
        self.output_ids = list(range(output_tokens))
        self.cached_tokens = 0
        self.dp_rank = -1

    @property
    def seqlen(self) -> int:
        return len(self.origin_input_ids) + len(self.output_ids)


class FakeScheduler(SchedulerMetricsMixin):
    def __init__(self):
        self.is_hybrid_swa = False
        self.is_hybrid_ssm = False
        self.waiting_queue = []
        self.disaggregation_mode = DisaggregationMode.DECODE
        self.disagg_decode_prealloc_queue = SimpleNamespace(
            pending_reqs=[], queue=[], retracted_queue=[]
        )
        self.disagg_decode_transfer_queue = SimpleNamespace(queue=[])
        self.running_batch = SimpleNamespace(reqs=[])
        self.dp_rank = 3

    def _get_token_info(self):
        return (100, None, None)


class TestSchedulerLoadAccounting(CustomTestCase):
    def test_decode_load_counts_pending_and_wrapped_handoff_reqs(self):
        scheduler = FakeScheduler()
        scheduler.running_batch.reqs = [FakeReq("running", 11)]
        scheduler.waiting_queue = [FakeReq("waiting", 7)]
        scheduler.disagg_decode_prealloc_queue.pending_reqs = [
            FakeReq("pending", 101)
        ]
        scheduler.disagg_decode_prealloc_queue.queue = [
            DecodeRequest(req=FakeReq("prealloc", 103), kv_receiver=None)
        ]
        scheduler.disagg_decode_transfer_queue.queue = [
            DecodeRequest(req=FakeReq("transfer", 107), kv_receiver=None)
        ]
        scheduler.disagg_decode_prealloc_queue.retracted_queue = [
            FakeReq("retracted", 109)
        ]

        load = scheduler.get_load()

        self.assertEqual(load.dp_rank, 3)
        self.assertEqual(load.pool_tokens, 100)
        self.assertEqual(load.num_waiting_reqs, 5)
        self.assertEqual(load.num_reqs, 6)
        self.assertEqual(load.num_tokens, 100 + 7 + 101 + 103 + 107 + 109)
        self.assertEqual(
            {info.rid: info.num_tokens for info in load.reqs_list},
            {
                "waiting": 7,
                "pending": 101,
                "prealloc": 103,
                "transfer": 107,
                "retracted": 109,
            },
        )


if __name__ == "__main__":
    import unittest

    unittest.main()