import unittest

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,
    VERY_LONG_PROMPT_LEN,
    base_engine_kwargs,
    run_until,
    run_until_all_finished,
    run_until_finished,
)


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

    def test_kv_pages_zero_after_finish(self):
        self.server.execute_script(self._script_kv_pages_zero_after_finish)

    @staticmethod
    def _script_kv_pages_zero_after_finish(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r)
        assert r.kv_pages == 0

    def test_kv_pages_positive_continuously_mid_chunk(self):
        self.server.execute_script(
            self._script_kv_pages_positive_continuously_mid_chunk
        )

    @staticmethod
    def _script_kv_pages_positive_continuously_mid_chunk(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)
        observed_chunking = False
        for _ in range(DEFAULT_MAX_STEPS):
            if r.is_chunking:
                observed_chunking = True
                assert (
                    r.kv_pages > 0
                ), f"kv_pages must be > 0 while is_chunking; got {r.kv_pages}"
            if r.finished:
                break
            yield
        assert observed_chunking, "test must observe at least one mid-chunk iter"
        assert r.finished

    def test_batch_composition_consistent_with_status(self):
        self.server.execute_script(
            self._script_batch_composition_consistent_with_status
        )

    @staticmethod
    def _script_batch_composition_consistent_with_status(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        for _ in range(DEFAULT_MAX_STEPS):
            if r.status == "running":
                comp = t.batch_composition()
                all_rids = (
                    comp.get("prefill", [])
                    + comp.get("decode", [])
                    + comp.get("chunked", [])
                )
                assert (
                    r.rid in all_rids
                ), f"running but not in batch_composition: {comp}"
            if r.finished:
                return
            yield
        raise AssertionError("req never finished")

    def test_active_reqs_listing(self):
        self.server.execute_script(self._script_active_reqs_listing)

    @staticmethod
    def _script_active_reqs_listing(t: ScriptedContext):
        r1 = t.start_req(prompt_len=16, max_new_tokens=4)
        r2 = t.start_req(prompt_len=16, max_new_tokens=4)
        yield
        actives = t.list_active_reqs()
        rids = {h.rid for h in actives}
        assert r1.rid in rids or r2.rid in rids
        yield from run_until_all_finished([r1, r2])
        for _ in range(12):
            actives_after = t.list_active_reqs()
            if all(h.rid not in (r1.rid, r2.rid) for h in actives_after):
                break
            yield
        actives_after = t.list_active_reqs()
        assert all(h.rid not in (r1.rid, r2.rid) for h in actives_after)

    def test_finished_means_chunks_done_stable(self):
        self.server.execute_script(self._script_finished_means_chunks_done_stable)

    @staticmethod
    def _script_finished_means_chunks_done_stable(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r)
        snap = r.chunks_done
        for _ in range(10):
            yield
            assert r.chunks_done == snap

    def test_finished_means_kv_pages_stays_zero(self):
        self.server.execute_script(self._script_finished_means_kv_pages_stays_zero)

    @staticmethod
    def _script_finished_means_kv_pages_stays_zero(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield from run_until_finished(r)
        for _ in range(10):
            yield
            assert r.kv_pages == 0

    def test_engine_stats_keys_present(self):
        self.server.execute_script(self._script_engine_stats_keys_present)

    @staticmethod
    def _script_engine_stats_keys_present(t: ScriptedContext):
        stats = t.engine_stats()
        assert isinstance(stats, dict)
        assert "kv_pool_free" in stats
        assert "req_pool_free" in stats
        yield

    def test_kv_pool_recovers_to_baseline(self):
        self.server.execute_script(self._script_kv_pool_recovers_to_baseline)

    @staticmethod
    def _script_kv_pool_recovers_to_baseline(t: ScriptedContext):
        before = t.engine_stats()["kv_pool_free"]
        reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(8)]
        yield from run_until_all_finished(reqs)
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        after = t.engine_stats()["kv_pool_free"]
        assert after >= before

    def test_hundred_reqs_no_leak(self):
        self.server.execute_script(self._script_hundred_reqs_no_leak)

    @staticmethod
    def _script_hundred_reqs_no_leak(t: ScriptedContext):
        baseline = t.engine_stats()
        reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(100)]
        yield from run_until_all_finished(reqs, max_steps=4000)
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        final = t.engine_stats()
        assert (
            final["kv_pool_free"] >= baseline["kv_pool_free"]
        ), f"KV leak: {baseline['kv_pool_free']} -> {final['kv_pool_free']}"
        assert final["req_pool_free"] >= baseline["req_pool_free"]

    def test_long_lived_engine_reps_chunked(self):
        self.server.execute_script(self._script_long_lived_engine_reps_chunked)

    @staticmethod
    def _script_long_lived_engine_reps_chunked(t: ScriptedContext):
        baseline = t.engine_stats()
        for _ in range(20):
            reqs = [
                t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
                for _ in range(5)
            ]
            yield from run_until_all_finished(reqs, max_steps=2000)
            for r in reqs:
                assert r.finished
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        final = t.engine_stats()
        assert final["kv_pool_free"] >= baseline["kv_pool_free"]

    def test_sustained_long_chunked_load(self):
        self.server.execute_script(self._script_sustained_long_chunked_load)

    @staticmethod
    def _script_sustained_long_chunked_load(t: ScriptedContext):
        expected_chunks_done = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE
        baseline_kv = t.engine_stats()["kv_pool_free"]
        reqs = [
            t.start_req(
                prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=10 + i
            )
            for i in range(30)
        ]
        yield from run_until_all_finished(reqs, max_steps=DEFAULT_MAX_STEPS * 20)
        for r in reqs:
            assert r.finished
            assert r.chunks_done == expected_chunks_done, (
                f"VERY_LONG_PROMPT_LEN must take exactly {expected_chunks_done} "
                f"chunks; got chunks_done={r.chunks_done}"
            )
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        final_kv = t.engine_stats()["kv_pool_free"]
        assert (
            final_kv >= baseline_kv
        ), f"KV leak after sustained chunked load: {baseline_kv} -> {final_kv}"

    def test_round_robin_short_and_chunked(self):
        self.server.execute_script(self._script_round_robin_short_and_chunked)

    @staticmethod
    def _script_round_robin_short_and_chunked(t: ScriptedContext):
        baseline = t.engine_stats()
        for _ in range(5):
            shorts = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(10)]
            chunked = [
                t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
                for _ in range(1)
            ]
            yield from run_until_all_finished(shorts + chunked, max_steps=2000)
            for r in shorts + chunked:
                assert r.finished
        for _ in range(40):
            if t.is_fully_idle:
                break
            yield
        t.flush_cache()
        yield
        final = t.engine_stats()
        assert final["kv_pool_free"] >= baseline["kv_pool_free"]

    def test_long_decode_then_many_short(self):
        self.server.execute_script(self._script_long_decode_then_many_short)

    @staticmethod
    def _script_long_decode_then_many_short(t: ScriptedContext):
        expected_chunks_done = VERY_LONG_PROMPT_LEN // DEFAULT_CHUNK_SIZE
        long_max_new_tokens = 256
        long_decode = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN,
            max_new_tokens=long_max_new_tokens,
            ignore_eos=True,
        )
        shorts = [t.start_req(prompt_len=8, max_new_tokens=2) for _ in range(50)]
        all_reqs = [long_decode] + shorts
        yield from run_until_all_finished(all_reqs, max_steps=DEFAULT_MAX_STEPS * 20)
        for r in all_reqs:
            assert r.finished
        assert long_decode.chunks_done == expected_chunks_done, (
            f"long req must chunk across exactly {expected_chunks_done} chunks; "
            f"got chunks_done={long_decode.chunks_done}"
        )
        assert len(long_decode.req.output_ids) == long_max_new_tokens, (
            f"ignore_eos long req must decode exactly {long_max_new_tokens} "
            f"tokens; got len(output_ids)={len(long_decode.req.output_ids)}"
        )

    def test_engine_stats_monotone_after_each_batch(self):
        self.server.execute_script(self._script_engine_stats_monotone_after_each_batch)

    @staticmethod
    def _script_engine_stats_monotone_after_each_batch(t: ScriptedContext):
        last = None
        for _ in range(10):
            reqs = [t.start_req(prompt_len=16, max_new_tokens=2) for _ in range(8)]
            yield from run_until_all_finished(reqs)
            for _ in range(40):
                if t.is_fully_idle:
                    break
                yield
            t.flush_cache()
            yield
            cur = t.engine_stats()["kv_pool_free"]
            if last is not None:
                assert cur >= last - 1, f"KV pool drifted: {last} -> {cur}"
            last = cur

    def test_inflight_middle_chunks_caps_at_one(self):
        self.server.execute_script(self._script_inflight_middle_chunks_caps_at_one)

    @staticmethod
    def _script_inflight_middle_chunks_caps_at_one(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        running_max = 0
        running_max_post_finish = 0
        post_finish_samples = 0
        for _ in range(DEFAULT_MAX_STEPS):
            yield
            req = r.req
            cur = req.inflight_middle_chunks if req is not None else 0
            running_max = max(running_max, cur)
            if r.finished:
                running_max_post_finish = max(running_max_post_finish, cur)
                post_finish_samples += 1
                if post_finish_samples >= 5:
                    break
        assert r.finished, "req never finished"
        assert running_max == 1, (
            f"inflight_middle_chunks must reach exactly 1 across the chunked "
            f"lifecycle (the cap from revert e875cd36e4); observed max={running_max}"
        )
        assert running_max_post_finish == 0, (
            f"inflight_middle_chunks must be reset to 0 after finish; "
            f"observed max post-finish={running_max_post_finish}"
        )

    def test_chunks_done_strictly_increases_no_plateaus(self):
        self.server.execute_script(
            self._script_chunks_done_strictly_increases_no_plateaus
        )

    @staticmethod
    def _script_chunks_done_strictly_increases_no_plateaus(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)
        prev_chunks_done = r.chunks_done
        prev_was_chunking = r.is_chunking
        for _ in range(DEFAULT_MAX_STEPS):
            yield
            if r.finished:
                return
            cur_chunks_done = r.chunks_done
            cur_is_chunking = r.is_chunking
            if prev_was_chunking and cur_is_chunking:
                assert cur_chunks_done > prev_chunks_done, (
                    f"chunks_done plateau between consecutive mid-chunk "
                    f"yields: {prev_chunks_done} -> {cur_chunks_done}"
                )
            prev_chunks_done = cur_chunks_done
            prev_was_chunking = cur_is_chunking
        raise AssertionError("req never finished")

    def test_output_tokens_len_equals_max_new_tokens_chunked(self):
        self.server.execute_script(
            self._script_output_tokens_len_equals_max_new_tokens_chunked
        )

    @staticmethod
    def _script_output_tokens_len_equals_max_new_tokens_chunked(t: ScriptedContext):
        n: int = 8
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN,
            max_new_tokens=n,
            ignore_eos=True,
        )
        yield from run_until_finished(r)
        assert r.finished
        assert (
            r.chunks_done >= 2
        ), f"VERY_LONG_PROMPT_LEN should chunk; got chunks_done={r.chunks_done}"
        assert len(r.req.output_ids) == n, (
            f"ignore_eos=True + max_new_tokens={n} must produce exactly "
            f"{n} output tokens; got len(output_tokens)={len(r.req.output_ids)}"
        )

    def test_num_input_tokens_equals_prompt_len_for_chunked(self):
        self.server.execute_script(
            self._script_num_input_tokens_equals_prompt_len_for_chunked
        )

    @staticmethod
    def _script_num_input_tokens_equals_prompt_len_for_chunked(t: ScriptedContext):
        prompt_len: int = VERY_LONG_PROMPT_LEN
        r = t.start_req(prompt_len=prompt_len, max_new_tokens=2)
        yield from run_until_finished(r)
        assert r.finished
        assert r.chunks_done >= 2, (
            f"VERY_LONG_PROMPT_LEN should chunk so this invariant is exercised "
            f"on a real multi-chunk prefill; got chunks_done={r.chunks_done}"
        )
        assert r.remaining_prompt_tokens == 0, (
            f"the whole prompt must be committed after a chunked finish; "
            f"remaining_prompt_tokens={r.remaining_prompt_tokens}"
        )

    def test_chunked_in_flight_count_exactly_zero_after_finish(self):
        self.server.execute_script(
            self._script_chunked_in_flight_count_exactly_zero_after_finish
        )

    @staticmethod
    def _script_chunked_in_flight_count_exactly_zero_after_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)
        assert (1 if t.scheduler.chunked_req is not None else 0) == 1, (
            f"chunked_in_flight_count should be 1 mid-chunk; got "
            f"{(1 if t.scheduler.chunked_req is not None else 0)}"
        )
        yield from run_until_finished(r)
        for _ in range(3):
            yield
            assert (1 if t.scheduler.chunked_req is not None else 0) == 0, (
                f"chunked_in_flight_count must be 0 after finish; got "
                f"{(1 if t.scheduler.chunked_req is not None else 0)}"
            )

    def test_extend_batch_idx_monotonic_invariant(self):
        self.server.execute_script(self._script_extend_batch_idx_monotonic_invariant)

    @staticmethod
    def _script_extend_batch_idx_monotonic_invariant(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=64)
        observed_regression: bool = False

        yield from run_until(
            r,
            lambda h: (
                h.req is not None
                and not h.req.is_retracted
                and h.req.extend_batch_idx > 0
                and h.remaining_prompt_tokens == 0
                and not h.finished
            ),
        )
        pre_retract_idx: int = r.req.extend_batch_idx
        assert not r.req.is_retracted

        t.pause_generation(mode="retract")
        retracted = t.find_req_by_rid(r.rid)
        assert retracted is not None, "retracted req must stay live in the queue"
        if retracted.extend_batch_idx < pre_retract_idx:
            observed_regression = True
            assert retracted.is_retracted, (
                f"extend_batch_idx regressed without retract flag: "
                f"{pre_retract_idx} -> {retracted.extend_batch_idx}"
            )
        t.continue_generation()

        prev_extend_batch_idx: int = -1
        regressions: int = 0
        for _ in range(DEFAULT_MAX_STEPS):
            req = t.find_req_by_rid(r.rid)
            if req is not None:
                cur_extend_batch_idx = req.extend_batch_idx
                if (
                    prev_extend_batch_idx >= 0
                    and cur_extend_batch_idx < prev_extend_batch_idx
                ):
                    regressions += 1
                    observed_regression = True
                    assert regressions == 1, (
                        f"extend_batch_idx regressed more than once for a single "
                        f"retract episode: "
                        f"{prev_extend_batch_idx} -> {cur_extend_batch_idx}"
                    )
                prev_extend_batch_idx = cur_extend_batch_idx
            if r.finished:
                break
            yield
        assert r.finished, "req never finished"
        assert observed_regression, (
            "retract must reset extend_batch_idx, producing the regression this "
            "test guards"
        )

    def test_inflight_decrement_only_on_final_invariant(self):
        self.server.execute_script(
            self._script_inflight_decrement_only_on_final_invariant
        )

    @staticmethod
    def _script_inflight_decrement_only_on_final_invariant(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        prev_inflight: int = 0
        prev_was_chunked_slot: bool = False
        prev_finished: bool = False
        observed_decrement: bool = False
        for _ in range(DEFAULT_MAX_STEPS):
            s = t.scheduler
            req = t.find_req_by_rid(r.rid)
            cur_inflight = req.inflight_middle_chunks if req is not None else 0
            cur_is_chunked_slot = (
                s.chunked_req is not None and s.chunked_req.rid == r.rid
            )
            cur_finished = req.finished() if req is not None else True
            if cur_inflight < prev_inflight:
                observed_decrement = True
                slot_just_released = prev_was_chunked_slot and not cur_is_chunked_slot
                finish_just_happened = (not prev_finished) and cur_finished
                assert slot_just_released or finish_just_happened, (
                    f"inflight_middle_chunks decreased ({prev_inflight} -> "
                    f"{cur_inflight}) without chunked slot release or req "
                    f"finish; prev_was_chunked_slot={prev_was_chunked_slot}, "
                    f"cur_is_chunked_slot={cur_is_chunked_slot}, "
                    f"prev_finished={prev_finished}, cur_finished={cur_finished}"
                )
            prev_inflight = cur_inflight
            prev_was_chunked_slot = cur_is_chunked_slot
            prev_finished = cur_finished
            if r.finished:
                break
            yield
        assert r.finished
        assert observed_decrement, (
            "test must observe at least one inflight_middle_chunks decrement "
            "across the chunked lifecycle"
        )


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