import unittest
from typing import Optional

from sglang.test.scripted_runtime.context import ScriptedContext
from sglang.test.scripted_runtime.test_case import ScriptedTestCase
from sglang.test.scripted_runtime_chunked_helpers import (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_MAX_STEPS,
    SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS,
    SMALL_KV_POOL_BALLAST_PROMPT_LEN,
    SMALL_KV_POOL_MAX_TOTAL_TOKENS,
    VERY_LONG_PROMPT_LEN,
    advance_to_decode_step,
    base_engine_kwargs,
    exhaust_row_pool,
    run_until,
    run_until_finished,
)


def _load_inquirer_pending_for_rid(t: ScriptedContext, rid: str) -> int:
    s = t.scheduler
    chunked = s.chunked_req
    if chunked is not None and chunked.rid == rid:
        return chunked.seqlen - len(chunked.prefix_indices)
    for req in s.waiting_queue:
        if req.rid == rid:
            return req.seqlen
    return 0


class TestSpecialCaseBasic(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_chunked_in_flight_no_idle(self):
        self.server.execute_script(self._script_chunked_in_flight_no_idle)

    @staticmethod
    def _script_chunked_in_flight_no_idle(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100
        )
        yield from run_until(r, lambda h: h.is_chunking)
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                saw_chunking = True
                assert (
                    not t.is_idle
                ), "scheduler must not idle while chunked_req is in flight"
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_chunking, "test must observe r mid-chunk at least once"

    def test_add_chunked_req_path(self):
        self.server.execute_script(self._script_add_chunked_req_path)

    @staticmethod
    def _script_add_chunked_req_path(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=110
        )
        yield from run_until_finished(r)
        assert r.finished
        assert r.chunks_done >= 2

    def test_admission_with_chunked_in_flight(self):
        self.server.execute_script(self._script_admission_with_chunked_in_flight)

    @staticmethod
    def _script_admission_with_chunked_in_flight(t: ScriptedContext):
        r_chunk = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=120
        )
        yield from run_until(r_chunk, lambda h: h.is_chunking)

        r_small = t.start_req(prompt_len=4, max_new_tokens=2)
        yield

        comp = t.batch_composition()
        assert r_chunk.rid in comp.get("chunked", [])

        yield from run_until_finished(r_small)
        yield from run_until_finished(r_chunk)
        assert r_small.finished and r_chunk.finished

    def test_abort_excludes_chunked_req(self):
        self.server.execute_script(self._script_abort_excludes_chunked_req)

    @staticmethod
    def _script_abort_excludes_chunked_req(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=130
        )
        yield from run_until(r, lambda h: h.is_chunking)

        t.abort(r)
        for _ in range(12):
            if t.scheduler.chunked_req is None and r.kv_pages == 0 and r.lock_refs == 0:
                break
            yield

        assert (
            t.scheduler.chunked_req is None
        ), f"abort must clear the chunked slot; got {t.scheduler.chunked_req!r}"
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_get_chunked_req_lambda_getter(self):
        self.server.execute_script(self._script_get_chunked_req_lambda_getter)

    @staticmethod
    def _script_get_chunked_req_lambda_getter(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=140
        )
        yield from run_until(r, lambda h: h.is_chunking)
        saw_match = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                cur = (
                    t.scheduler.chunked_req.rid
                    if t.scheduler.chunked_req is not None
                    else None
                )
                assert cur in (None, r.rid), (
                    f"getter returned unrelated rid: got {cur!r}, expected "
                    f"None or {r.rid!r}"
                )
                if cur == r.rid:
                    saw_match = True
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_match, "getter must return r.rid at least once while r.is_chunking"
        assert (
            t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None
        ) is None

    @unittest.skip(
        "requires real disaggregation prefill/decode split — single-engine "
        "ScriptedContext cannot exercise the decode-side waiting_queue "
        "KV-hold path. Belongs in test_scripted_disagg.py with D3 counter "
        "wiring once disagg topology is available."
    )
    def test_disagg_decode_waiting_queue_kv_held(self):
        pass

    @unittest.skip(
        "requires DLLM model + staging mixin — single-engine cannot drive "
        "both DLLM staging AND chunked admission incrementing "
        "inflight_middle_chunks from two sources. Belongs in DLLM-specific "
        "test file."
    )
    def test_dllm_staging_double_inflight_middle_chunks(self):
        pass

    @unittest.skip(
        "requires real disaggregation topology — single-engine cannot "
        "exercise staging_handler chunked path."
    )
    def test_staging_handler_chunked(self):
        pass

    @unittest.skip(
        "requires mooncake KV transport backend — single-engine cannot "
        "drive the conn layer chunked path."
    )
    def test_mooncake_conn_chunked(self):
        pass

    @unittest.skip(
        "requires NIXL KV transport backend — single-engine cannot drive "
        "the NIXL conn layer chunked path."
    )
    def test_nixl_conn_chunked(self):
        pass

    def test_filter_batch_exclude_chunked_flag(self):
        self.server.execute_script(self._script_filter_batch_exclude_chunked_flag)

    @staticmethod
    def _script_filter_batch_exclude_chunked_flag(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=150
        )
        r2 = t.start_req(prompt_len=16, max_new_tokens=2)
        saw_r1_chunking = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            if r1.is_chunking:
                saw_r1_chunking = True
                comp = t.batch_composition()
                assert r1.rid in comp.get(
                    "chunked", []
                ), f"mid-chunk r1 must occupy the chunked role; got {comp!r}"
                assert r1.rid not in comp.get(
                    "running", []
                ), f"chunked r1 must be excluded from the running role; got {comp!r}"
            if r1.finished and r2.finished:
                break
            yield
        assert r1.finished and r2.finished
        assert (
            saw_r1_chunking
        ), "r1 must have chunked at some point to exercise the exclude branch"

    @unittest.skip(
        "pdmux split_prefill_batch requires the pdmux topology — "
        "single-engine ScriptedContext cannot drive the split path. "
        "Belongs in a pdmux-specific test once that lane is wired up."
    )
    def test_pdmux_split_prefill_batch(self):
        pass

    def test_streaming_session_kv_committed_bound(self):
        self.server.execute_script(self._script_streaming_session_kv_committed_bound)

    @staticmethod
    def _script_streaming_session_kv_committed_bound(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=160
        )
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                assert len(r.req.prefix_indices) <= r.req.kv_committed_len, (
                    f"streaming-session chunked stash must stay bounded by "
                    f"kv_committed_len; prefix_indices_len={len(r.req.prefix_indices)}, "
                    f"kv_committed_len={r.req.kv_committed_len}"
                )
            if r.finished:
                break
            yield
        assert r.finished

    @unittest.skip(
        "mamba_pool_idx cleanup applies only to mamba-class models — "
        "single-engine with a non-mamba model cannot drive the NO_TOKEN "
        "chunked-resume cleanup-skip branch. Belongs in a mamba-specific "
        "test file once mamba scripted coverage is added."
    )
    def test_mamba_pool_idx_cleanup_skip_chunked_resume(self):
        pass

    def test_pause_retract_clears_chunked_req(self):
        self.server.execute_script(self._script_pause_retract_clears_chunked_req)

    @staticmethod
    def _script_pause_retract_clears_chunked_req(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=170
        )
        yield from run_until(r, lambda h: h.is_chunking)

        t.pause_generation(mode="retract")
        yield

        assert (
            t.scheduler.chunked_req is None
        ), f"pause(retract) must clear chunked_req; got {t.scheduler.chunked_req!r}"
        assert not r.finished, "retract must re-queue r, not finish or abort it"
        assert r.status == "waiting", (
            f"retracted chunked req must return to the waiting queue; "
            f"got status={r.status!r}"
        )

        t.continue_generation()
        yield from run_until_finished(r)
        assert (
            r.finished
        ), "continue_generation must drive the re-queued req to completion"

    def test_retract_during_gap_inflight_middle_chunks_positive(self):
        self.server.execute_script(
            self._script_retract_during_gap_inflight_middle_chunks_positive
        )

    @staticmethod
    def _script_retract_during_gap_inflight_middle_chunks_positive(t: ScriptedContext):
        r = t.start_req(
            prompt_len=3 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=180
        )
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
        assert r.req.inflight_middle_chunks > 0
        assert r.is_chunking

        t.pause_generation(mode="retract")
        yield

        assert r.kv_pages == 0, f"retract must release KV; got kv_pages={r.kv_pages}"
        assert not r.finished, "retract must re-queue r, not finish or abort it"
        assert not r.is_chunking, "retract must release the chunked slot"
        assert r.status == "waiting", (
            f"retracted chunked req must return to the waiting queue; "
            f"got status={r.status!r}"
        )
        req = t.find_req_by_rid(r.rid)
        assert req is not None and req.inflight_middle_chunks == 0, (
            f"retract must reset inflight_middle_chunks; got "
            f"{req.inflight_middle_chunks if req is not None else None}"
        )

        t.continue_generation()
        yield from run_until_finished(r, max_steps=2000)
        assert (
            r.finished
        ), "continue_generation must drive the re-queued req to completion"
        assert r.kv_pages == 0
        assert len(r.req.output_ids) == 2

    def test_load_inquirer_pending_tokens_dedup_chunked(self):
        self.server.execute_script(
            self._script_load_inquirer_pending_tokens_dedup_chunked
        )

    @staticmethod
    def _script_load_inquirer_pending_tokens_dedup_chunked(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=190
        )
        yield from run_until(r, lambda h: h.is_chunking)
        saw_chunking = False
        saw_dedup = False
        for _ in range(DEFAULT_MAX_STEPS):
            chunked = t.scheduler.chunked_req
            if r.is_chunking and chunked is not None and chunked.rid == r.rid:
                saw_chunking = True
                pending = _load_inquirer_pending_for_rid(t, r.rid)
                expected = chunked.seqlen - len(chunked.prefix_indices)
                assert pending == expected, (
                    f"load_inquirer chunked contribution must equal the prefix-"
                    f"subtracting formula seqlen - len(prefix_indices) = {expected}; "
                    f"got {pending}"
                )
                assert pending <= chunked.seqlen, (
                    f"chunked contribution must never exceed its full seqlen "
                    f"{chunked.seqlen}; got {pending} — dual-queue dedup violated"
                )
                if len(chunked.prefix_indices) > 0:
                    saw_dedup = pending < chunked.seqlen
            if r.finished:
                break
            yield
        assert r.finished
        assert (
            saw_chunking
        ), "test must observe the dual-queue chunked state at least once"
        assert saw_dedup, (
            "test must observe the chunked req with a committed prefix so the "
            "dedup subtraction is actually exercised"
        )

    def test_load_inquirer_chunked_contribution_exact_remainder(self):
        self.server.execute_script(
            self._script_load_inquirer_chunked_contribution_exact_remainder
        )

    @staticmethod
    def _script_load_inquirer_chunked_contribution_exact_remainder(t: ScriptedContext):
        s = t.scheduler
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=200
        )
        yield from run_until(r, lambda h: h.is_chunking)
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            chunked = s.chunked_req
            if r.is_chunking and chunked is not None and chunked.rid == r.rid:
                assert len(s.waiting_queue) == 0, (
                    "test requires an empty waiting_queue so the chunked req is "
                    f"the sole pending-token contributor; got {len(s.waiting_queue)}"
                )
                expected = chunked.seqlen - len(chunked.prefix_indices)
                observed = s.load_inquirer._get_num_pending_tokens()
                assert observed == expected, (
                    f"chunked contribution must equal remainder "
                    f"seqlen - len(prefix_indices) = {expected}, got {observed}; "
                    f"a value of {chunked.seqlen} would mean the committed prefix "
                    "is being double-counted"
                )
                saw_chunking = True
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_chunking, "test must observe the chunked req mid-flight"

    def test_load_inquirer_chunk_deduct_subtracts_planned_chunk(self):
        self.server.execute_script(
            self._script_load_inquirer_chunk_deduct_subtracts_planned_chunk
        )

    @staticmethod
    def _script_load_inquirer_chunk_deduct_subtracts_planned_chunk(t: ScriptedContext):
        s = t.scheduler
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=210
        )
        yield from run_until(r, lambda h: h.is_chunking)
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            chunked = s.chunked_req
            if (
                r.is_chunking
                and chunked is not None
                and chunked.rid == r.rid
                and chunked.extend_range.length > 0
            ):
                deduct = chunked.extend_range.length
                base = s.load_inquirer._get_num_pending_tokens()
                deducted = s.load_inquirer._get_num_pending_tokens(chunk_deduct=deduct)
                assert deducted == base - deduct, (
                    f"chunk_deduct must subtract the planned chunk exactly: "
                    f"base={base}, deducted={deducted}, expected={base - deduct}"
                )
                saw_chunking = True
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_chunking, "test must observe the chunked req with a planned chunk"

    def test_stage_a_inflight_middle_chunks_sync_invariant(self):
        self.server.execute_script(
            self._script_stage_a_inflight_middle_chunks_sync_invariant
        )

    @staticmethod
    def _script_stage_a_inflight_middle_chunks_sync_invariant(t: ScriptedContext):
        def assert_invariant() -> None:
            req = t.find_req_by_rid(r.rid)
            if req is not None and req.inflight_middle_chunks > 0:
                assert r.is_chunking, (
                    f"invariant violated: inflight_middle_chunks="
                    f"{req.inflight_middle_chunks} but is_chunking={r.is_chunking}"
                )

        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, prompt_token=220
        )
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
        assert_invariant()

        t.pause_generation(mode="retract")
        yield
        assert_invariant()

        t.continue_generation()
        for _ in range(DEFAULT_MAX_STEPS):
            assert_invariant()
            if r.finished:
                return
            yield
        raise AssertionError("req did not finish")

    def test_init_next_round_input_resets_chunk_state(self):
        self.server.execute_script(
            self._script_init_next_round_input_resets_chunk_state
        )

    @staticmethod
    def _script_init_next_round_input_resets_chunk_state(t: ScriptedContext):
        r = t.start_req(
            prompt_len=3 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=230
        )
        yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking)
        saw_mid_chunk = False
        for _ in range(DEFAULT_MAX_STEPS):
            req = r.req
            if (
                r.is_chunking
                and r.chunks_done >= 1
                and req is not None
                and req.extend_range is not None
            ):
                saw_mid_chunk = True
                assert (
                    req.extend_range.end
                    == len(req.prefix_indices) + req.extend_range.length
                ), (
                    f"init_next_round_input must rebuild fill_ids to the committed "
                    f"prefix plus the in-flight chunk; "
                    f"fill_ids_len={req.extend_range.end}, "
                    f"prefix_indices_len={len(req.prefix_indices)}, "
                    f"extend_input_len={req.extend_range.length}, "
                    f"chunks_done={r.chunks_done}"
                )
            if r.finished:
                break
            yield
        assert r.finished
        assert (
            saw_mid_chunk
        ), "test must observe the fill_ids reset boundary at least once"
        assert r.finished

    def test_chunked_req_slot_cleared_when_chunk_completes(self):
        self.server.execute_script(
            self._script_chunked_req_slot_cleared_when_chunk_completes
        )

    @staticmethod
    def _script_chunked_req_slot_cleared_when_chunk_completes(t: ScriptedContext):
        s = t.scheduler
        r = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=240
        )
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                saw_chunking = True
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_chunking, "req should have occupied the chunked_req slot mid-chunk"
        assert (
            s.chunked_req is None
        ), f"chunked_req slot must clear after last chunk; got {s.chunked_req!r}"

    def test_second_chunked_admit_blocked_when_chunked_req_set(self):
        self.server.execute_script(
            self._script_second_chunked_admit_blocked_when_chunked_req_set
        )

    @staticmethod
    def _script_second_chunked_admit_blocked_when_chunked_req_set(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10
        )
        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=11
        )
        saw_r1_chunking = False
        saw_r2_chunking = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            both = r1.is_chunking and r2.is_chunking
            assert not both, (
                f"only one chunked-resume slot allowed; r1.is_chunking="
                f"{r1.is_chunking}, r2.is_chunking={r2.is_chunking}"
            )
            if r1.is_chunking:
                saw_r1_chunking = True
            if r2.is_chunking:
                saw_r2_chunking = True
            if r1.finished and r2.finished:
                break
            yield
        assert r1.finished and r2.finished
        assert saw_r1_chunking and saw_r2_chunking, (
            f"both reqs should chunk over their lifetime; saw_r1="
            f"{saw_r1_chunking}, saw_r2={saw_r2_chunking}"
        )

    def test_scheduler_continues_with_only_chunked_req_no_waiting(self):
        self.server.execute_script(
            self._script_scheduler_continues_with_only_chunked_req_no_waiting
        )

    @staticmethod
    def _script_scheduler_continues_with_only_chunked_req_no_waiting(
        t: ScriptedContext,
    ):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=250
        )
        yield from run_until(r, lambda h: h.is_chunking)
        prev_chunks_done = r.chunks_done
        progressed = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                assert (
                    not t.is_idle
                ), "scheduler must not go idle while a chunked req is in flight"
            cur_chunks_done = r.chunks_done
            if cur_chunks_done > prev_chunks_done:
                progressed = True
            prev_chunks_done = cur_chunks_done
            if r.finished:
                break
            yield
        assert r.finished
        assert progressed, (
            "chunks_done must keep advancing without any waiter; pre-fix "
            "an empty waiting_queue could cause the loop to skip continuation"
        )


class TestSpecialCaseRowPoolExhaustion(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        max_running_requests=8,
    )

    def test_chunked_req_bypasses_req_pool_exhaustion(self):
        self.server.execute_script(
            self._script_chunked_req_bypasses_req_pool_exhaustion
        )

    @staticmethod
    def _script_chunked_req_bypasses_req_pool_exhaustion(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking)
        chunks_before_pressure = r.chunks_done

        yield from exhaust_row_pool(t, leave_rows=0)

        progressed_under_pressure = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            if r.chunks_done > chunks_before_pressure:
                progressed_under_pressure = True
            if r.finished:
                break
            yield
        assert r.finished
        assert progressed_under_pressure, (
            "chunked req must advance even when get_num_allocatable_reqs "
            "returns 0; pre-fix the bypass would block forever"
        )
        assert r.kv_pages == 0


class TestSpecialCaseMixedChunk(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        enable_mixed_chunk=True,
    )

    def test_mix_with_running_chunked_plus_decode(self):
        self.server.execute_script(self._script_mix_with_running_chunked_plus_decode)

    @staticmethod
    def _script_mix_with_running_chunked_plus_decode(t: ScriptedContext):
        decodes = [t.start_req(prompt_len=8, max_new_tokens=16) for _ in range(3)]
        for d in decodes:
            yield from advance_to_decode_step(d, 1)

        r_chunk = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=4, prompt_token=300
        )
        yield from run_until(r_chunk, lambda h: h.is_chunking)

        comp = t.batch_composition()
        assert r_chunk.rid in comp.get("chunked", [])

        saw_mixed_with_decode = False
        all_reqs = [r_chunk, *decodes]
        for _ in range(DEFAULT_MAX_STEPS * 2):
            comp = t.batch_composition()
            batch_rids = (
                set(comp.get("prefill", []))
                | set(comp.get("decode", []))
                | set(comp.get("running", []))
            )
            if (
                t.last_batch_forward_mode == "MIXED"
                and r_chunk.rid in comp.get("chunked", [])
                and any(d.rid in batch_rids for d in decodes)
            ):
                saw_mixed_with_decode = True
            if all(x.finished for x in all_reqs):
                break
            yield
        assert all(x.finished for x in all_reqs)
        assert saw_mixed_with_decode, (
            "enable_mixed_chunk should merge running decode reqs into the chunked "
            "prefill iter (MIXED forward_mode) at least once"
        )

    def test_mixed_chunk_with_logprob_falls_back(self):
        self.server.execute_script(self._script_mixed_chunk_with_logprob_falls_back)

    @staticmethod
    def _script_mixed_chunk_with_logprob_falls_back(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN,
            max_new_tokens=2,
            return_logprob=True,
            prompt_token=310,
        )
        yield from run_until(r, lambda h: h.is_chunking)
        assert (
            t.last_batch_forward_mode != "MIXED"
        ), f"return_logprob must disable mixed-chunk path; got {t.last_batch_forward_mode!r}"
        yield from run_until_finished(r)

    def test_mixed_chunk_with_running_batch(self):
        self.server.execute_script(self._script_mixed_chunk_with_running_batch)

    @staticmethod
    def _script_mixed_chunk_with_running_batch(t: ScriptedContext):
        r_dec = t.start_req(prompt_len=8, max_new_tokens=32)
        yield from run_until(r_dec, lambda h: h.status == "running")

        r_chunk = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=320
        )
        yield
        yield from run_until(r_chunk, lambda h: h.is_chunking)

        assert (
            t.last_batch_forward_mode == "MIXED"
        ), f"chunked admission with running batch must enter MIXED; got {t.last_batch_forward_mode!r}"
        for _ in range(DEFAULT_MAX_STEPS * 2):
            if r_chunk.finished and r_dec.finished:
                break
            yield
        assert r_chunk.finished and r_dec.finished


class TestSpecialCaseTransformers(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        model_impl="transformers",
    )

    def test_transformers_text_model_still_chunks(self):
        self.server.execute_script(self._script_transformers_text_model_still_chunks)

    @staticmethod
    def _script_transformers_text_model_still_chunks(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10
        )
        yield from run_until_finished(r)
        assert r.finished
        expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE)
        assert r.chunks_done == expected, (
            f"a text model on the Transformers backend must chunk like the native "
            f"backend ({expected} chunks); got chunks_done={r.chunks_done}"
        )


class TestSpecialCaseNoChunking(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=-1)

    def test_chunk_size_negative_disables_chunking(self):
        self.server.execute_script(self._script_chunk_size_negative_disables_chunking)

    @staticmethod
    def _script_chunk_size_negative_disables_chunking(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        for _ in range(DEFAULT_MAX_STEPS):
            assert (
                not r.is_chunking
            ), "chunked_prefill_size=-1 should disable chunked path"
            if r.finished:
                return
            yield
        raise AssertionError("req did not finish under disabled chunking")


DETERMINISTIC_ALIGN_SIZE = 4096


class TestSpecialCaseDeterministicFlashInfer(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DETERMINISTIC_ALIGN_SIZE,
        page_size=16,
        attention_backend="flashinfer",
        enable_deterministic_inference=True,
        kv_canary="none",
        kv_canary_sweep_interval=0,
    )

    def test_chunked_truncation_align_size(self):
        self.server.execute_script(self._script_chunked_truncation_align_size)

    @staticmethod
    def _script_chunked_truncation_align_size(t: ScriptedContext):
        r = t.start_req(
            prompt_len=DETERMINISTIC_ALIGN_SIZE + 1024,
            max_new_tokens=2,
            prompt_token=10,
        )
        page_size = 16
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking and r.req.extend_range is not None:
                saw_chunking = True
                assert r.req.extend_range.length % page_size == 0, (
                    f"deterministic chunk boundary must be page-aligned; "
                    f"got extend_input_len={r.req.extend_range.length}, page_size={page_size}"
                )
            if r.finished:
                break
            yield
        assert r.finished, "chunked req did not finish"
        assert saw_chunking, "test must observe the req mid-chunk at least once"


class TestSpecialCaseHiCache(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        enable_hierarchical_cache=True,
        kv_canary="none",
        kv_canary_sweep_interval=0,
    )

    def test_hicache_breakdown_only_first_chunk(self):
        self.server.execute_script(self._script_hicache_breakdown_only_first_chunk)

    @staticmethod
    def _script_hicache_breakdown_only_first_chunk(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=400
        )
        first_chunk_snap = None
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                saw_chunking = True
            if r.is_chunking and r.chunks_done >= 1 and first_chunk_snap is None:
                first_chunk_snap = r.req.cached_tokens
            if first_chunk_snap is not None and r.is_chunking:
                cur = r.req.cached_tokens
                assert cur == first_chunk_snap, (
                    f"HiCache cached_tokens_* must freeze after first chunk; "
                    f"first={first_chunk_snap!r}, now={cur!r}"
                )
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_chunking, "test must observe r mid-chunk at least once"
        assert (
            first_chunk_snap is not None
        ), "test must snapshot cached_tokens at the first chunk boundary"

    def test_hicache_cached_tokens_set_once_invariant(self):
        self.server.execute_script(
            self._script_hicache_cached_tokens_set_once_invariant
        )

    @staticmethod
    def _script_hicache_cached_tokens_set_once_invariant(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=410
        )
        snap: Optional[tuple] = None
        saw_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            req = t.find_req_by_rid(r.rid)
            if req is not None and r.is_chunking and r.chunks_done >= 1:
                saw_chunking = True
                cur = (
                    req.cached_tokens_device,
                    req.cached_tokens_host,
                    req.cached_tokens_storage,
                )
                if snap is None:
                    snap = cur
                else:
                    assert cur == snap, (
                        f"HiCache cached_tokens_* fields must be set exactly "
                        f"once on first chunk; values changed: snap={snap}, "
                        f"cur={cur}"
                    )
            if r.finished:
                break
            yield
        assert r.finished
        assert (
            saw_chunking
        ), "test must observe the req mid-chunk (chunks_done >= 1) at least once"
        assert snap is not None, "test must snapshot the cached_tokens_* breakdown"


def _expected_chunks(prompt_len: int, chunk_size: int) -> int:
    if prompt_len <= chunk_size:
        return 0
    return (prompt_len + chunk_size - 1) // chunk_size


class TestSpecialCaseDynamicChunkingPP1(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        enable_dynamic_chunking=True,
    )

    def test_dynamic_chunking_forced_off_on_pp1(self):
        self.server.execute_script(self._script_dynamic_chunking_forced_off_on_pp1)

    @staticmethod
    def _script_dynamic_chunking_forced_off_on_pp1(t: ScriptedContext):
        assert t.scheduler.enable_dynamic_chunking is False, (
            "pp_size==1 must force enable_dynamic_chunking off even when the "
            "server arg is True (the 'and ps.pp_size > 1' conjunct)"
        )
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10
        )
        yield from run_until_finished(r)
        assert r.finished
        expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE)
        assert r.chunks_done == expected, (
            f"uniform chunked_prefill_size must yield exactly {expected} chunks "
            f"(no dynamic-size prediction on pp1); got {r.chunks_done}"
        )


class TestSpecialCaseSmallPool(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        max_total_tokens=SMALL_KV_POOL_MAX_TOTAL_TOKENS,
    )

    @staticmethod
    def _start_ballast(t: ScriptedContext, *, prompt_token: int):
        return t.start_req(
            prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN,
            max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS,
            ignore_eos=True,
            prompt_token=prompt_token,
        )

    @staticmethod
    def _run_force_readd_then_complete(
        t: ScriptedContext, *, chunk_token: int, ballast_token: int
    ):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=chunk_token
        )
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)

        ballast = TestSpecialCaseSmallPool._start_ballast(t, prompt_token=ballast_token)

        ballast_retracted = False
        for _ in range(2000):
            if ballast.status == "waiting":
                ballast_retracted = True
            if r.finished:
                break
            yield
        assert r.finished, (
            "force-re-added chunked resume must complete once the engine retracts "
            f"the ballast; status={r.status} kv_pages={r.kv_pages}"
        )
        assert r.kv_pages == 0, f"kv_pages={r.kv_pages}"
        assert r.lock_refs == 0, f"lock_refs={r.lock_refs}"
        assert (
            ballast_retracted
            or ballast.finished
            or ballast.status in ("waiting", "finished", "unknown")
        ), f"ballast must be retracted/aborted under pressure; status={ballast.status}"

        t.abort(ballast)
        for _ in range(200):
            if t.is_fully_idle:
                break
            yield

    def test_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens(self):
        self.server.execute_script(
            self._script_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens
        )

    @staticmethod
    def _script_add_chunked_req_rem_nonpositive_forces_rem_chunk_tokens(
        t: ScriptedContext,
    ):
        yield from TestSpecialCaseSmallPool._run_force_readd_then_complete(
            t, chunk_token=700, ballast_token=701
        )

    def test_chunked_forced_admission_avoids_leak(self):
        self.server.execute_script(self._script_chunked_forced_admission_avoids_leak)

    @staticmethod
    def _script_chunked_forced_admission_avoids_leak(t: ScriptedContext):
        yield from TestSpecialCaseSmallPool._run_force_readd_then_complete(
            t, chunk_token=710, ballast_token=711
        )

    def test_add_chunked_req_non_swa_forced_admit_on_rem_zero(self):
        self.server.execute_script(
            self._script_add_chunked_req_non_swa_forced_admit_on_rem_zero
        )

    @staticmethod
    def _script_add_chunked_req_non_swa_forced_admit_on_rem_zero(t: ScriptedContext):
        yield from TestSpecialCaseSmallPool._run_force_readd_then_complete(
            t, chunk_token=720, ballast_token=721
        )


class TestSpecialCaseRetractMerge(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_retract_merges_extend_chunk_batch_before_retract_all(self):
        self.server.execute_script(
            self._script_retract_merges_extend_chunk_batch_before_retract_all
        )

    @staticmethod
    def _script_retract_merges_extend_chunk_batch_before_retract_all(
        t: ScriptedContext,
    ):
        s = t.scheduler
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking)

        assert s.last_batch is not None, "last_batch must be set while chunking"
        assert s.last_batch.forward_mode.is_extend(), (
            f"last_batch must be the extend/chunk batch; got "
            f"{s.last_batch.forward_mode!r}"
        )

        t.pause_generation(mode="retract")
        yield

        assert (
            s.last_batch is None
        ), "retract must clear last_batch after merging the extend chunk batch"
        assert len(s.running_batch.reqs) == 0, (
            "the merged extend chunk batch must be retracted out of running_batch, "
            f"not stranded; got {len(s.running_batch.reqs)} reqs"
        )
        assert (
            r.status == "waiting"
        ), f"retracted chunked req must return to the waiting queue; got {r.status!r}"
        assert r.kv_pages == 0

        t.continue_generation()
        yield from run_until_finished(r)
        assert r.finished
        assert len(r.req.output_ids) == 2, (
            f"resumed req must emit exactly max_new_tokens; got "
            f"{len(r.req.output_ids)}"
        )


class TestSpecialCaseChunkBudgetDefer(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_co_submitted_waiter_deferred_when_chunk_budget_zero(self):
        self.server.execute_script(
            self._script_co_submitted_waiter_deferred_when_chunk_budget_zero
        )

    @staticmethod
    def _script_co_submitted_waiter_deferred_when_chunk_budget_zero(
        t: ScriptedContext,
    ):
        r_chunk = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        r_wait = t.start_req(prompt_len=DEFAULT_CHUNK_SIZE + 8, max_new_tokens=2)
        yield from run_until(r_chunk, lambda h: h.is_chunking)

        saw_deferred = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            comp = t.batch_composition()
            chunk_active = r_chunk.rid in comp.get("chunked", [])
            wait_idle = (
                r_wait.status == "waiting"
                and r_wait.rid not in comp.get("chunked", [])
                and r_wait.rid not in comp.get("running", [])
                and r_wait.rid not in comp.get("decode", [])
            )
            if chunk_active and wait_idle:
                saw_deferred = True
                assert r_wait.kv_pages == 0, (
                    f"deferred waiter (rem_chunk_tokens<=0 -> OTHER) must hold no KV; "
                    f"got kv_pages={r_wait.kv_pages}"
                )
            if r_chunk.finished and r_wait.finished:
                break
            yield
        assert r_chunk.finished and r_wait.finished
        assert saw_deferred, (
            "second waiter must be deferred for >=1 iter while the chunk-token "
            "budget was driven to 0 by the in-flight chunked req"
        )


class TestSpecialCaseIgnoreEosNoRadix(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        disable_radix_cache=True,
        kv_canary="none",
        kv_canary_sweep_interval=0,
    )

    def test_ignore_eos_chunked_truncate_path(self):
        self.server.execute_script(self._script_ignore_eos_chunked_truncate_path)

    @staticmethod
    def _script_ignore_eos_chunked_truncate_path(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True
        )
        yield from run_until(r, lambda h: h.is_chunking)
        yield from run_until_finished(r)
        assert r.finished
        expected = _expected_chunks(VERY_LONG_PROMPT_LEN, DEFAULT_CHUNK_SIZE)
        assert r.chunks_done == expected, (
            f"ignore_eos chunk-truncation path must carry the prompt as new_chunked_req "
            f"for {expected} chunks; got {r.chunks_done}"
        )

    def test_ignore_eos_nonchunked_fits_chunk_budget(self):
        self.server.execute_script(self._script_ignore_eos_nonchunked_fits_chunk_budget)

    @staticmethod
    def _script_ignore_eos_nonchunked_fits_chunk_budget(t: ScriptedContext):
        r = t.start_req(prompt_len=8, max_new_tokens=2, ignore_eos=True)
        for _ in range(DEFAULT_MAX_STEPS):
            assert not r.is_chunking, (
                "a prompt that fits the chunk budget must admit non-chunked via the "
                "ignore_eos path, never entering is_chunking"
            )
            if r.finished:
                break
            yield
        assert r.finished
        assert r.chunks_done == 0, (
            f"prompt fitting the chunk budget completes non-chunked; got "
            f"chunks_done={r.chunks_done}"
        )

    def test_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero(self):
        self.server.execute_script(
            self._script_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero
        )

    @staticmethod
    def _script_ignore_eos_second_long_waiter_deferred_on_chunk_budget_zero(
        t: ScriptedContext,
    ):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True
        )
        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True
        )
        yield from run_until(r1, lambda h: h.is_chunking)

        saw_deferred = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            comp = t.batch_composition()
            r1_active = r1.rid in comp.get("chunked", [])
            r2_idle = (
                r2.status == "waiting"
                and r2.rid not in comp.get("chunked", [])
                and r2.rid not in comp.get("running", [])
                and r2.rid not in comp.get("decode", [])
            )
            if r1_active and r2_idle:
                saw_deferred = True
                assert r2.kv_pages == 0, (
                    f"deferred ignore_eos waiter must hold no KV; got "
                    f"kv_pages={r2.kv_pages}"
                )
            if r1.finished and r2.finished:
                break
            yield
        assert r1.finished and r2.finished
        assert saw_deferred, (
            "second long ignore_eos waiter must be deferred for >=1 iter while the "
            "chunk-token budget was driven to 0 (nested rem_chunk_tokens<=0 OTHER)"
        )


class TestSpecialCaseRetractedStain(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_retracted_stain_suppresses_cached_token_recount(self):
        self.server.execute_script(
            self._script_retracted_stain_suppresses_cached_token_recount
        )

    @staticmethod
    def _script_retracted_stain_suppresses_cached_token_recount(t: ScriptedContext):
        r = t.start_req(prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
        cached_before = r.req.cached_tokens

        t.pause_generation(mode="retract")
        yield
        t.continue_generation()
        yield from run_until_finished(r, max_steps=2000)
        assert r.finished

        req = r.req
        assert (
            req.retracted_stain is True
        ), "retract must set retracted_stain so the cached-token recount is suppressed"
        assert req.cached_tokens == cached_before, (
            f"retracted_stain must suppress re-adding pre_len-already_computed on "
            f"resume; cached_tokens grew from {cached_before} to {req.cached_tokens}"
        )


class TestSpecialCaseResultSkipRetracted(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_result_skip_retracted_emits_no_stray_token(self):
        self.server.execute_script(
            self._script_result_skip_retracted_emits_no_stray_token
        )

    @staticmethod
    def _script_result_skip_retracted_emits_no_stray_token(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, ignore_eos=True
        )
        yield from run_until(r, lambda h: h.is_chunking)

        t.pause_generation(mode="retract")
        yield
        t.continue_generation()
        yield from run_until_finished(r, max_steps=2000)
        assert r.finished
        assert len(r.req.output_ids) == 2, (
            f"the skipped retracted entry must not append a stray output token; "
            f"got output_ids len {len(r.req.output_ids)}"
        )


class TestSpecialCaseMiddleChunkNoToken(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_middle_chunk_appends_no_token_no_finish(self):
        self.server.execute_script(self._script_middle_chunk_appends_no_token_no_finish)

    @staticmethod
    def _script_middle_chunk_appends_no_token_no_finish(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until(r, lambda h: h.is_chunking)

        saw_middle_chunk = False
        for _ in range(DEFAULT_MAX_STEPS * 2):
            if r.is_chunking:
                saw_middle_chunk = True
                assert len(r.req.output_ids) == 0, (
                    f"middle chunk must not append an output token; got "
                    f"output_ids len {len(r.req.output_ids)}"
                )
                assert (
                    r.status != "finished"
                ), "middle chunk must not finish the req (skip_stream_req)"
            if r.finished:
                break
            yield
        assert r.finished
        assert saw_middle_chunk, "test must observe r mid-chunk at least once"
        assert (
            len(r.req.output_ids) >= 1
        ), "output tokens must appear only after the chunked prefill completes"


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