for i, running_req in enumerate(self.running_batch.reqs):
    if running_req in preemptible_reqs:
        self.rem_total_token_offset -= (
            self._get_running_request_total_token_offset(req)  # ❌ Wrong: should be running_req
        )

for i, running_req in enumerate(self.running_batch.reqs):
    if running_req in preemptible_reqs:
        self.rem_total_token_offset -= (
            self._get_running_request_total_token_offset(running_req)  # ✅ Correct
        )

"""
Unit test to reproduce the rem_total_token_offset bug in schedule_policy.py:700
This test demonstrates the incorrect variable usage during preemption.
"""

import unittest
from unittest.mock import Mock, MagicMock
from sglang.srt.managers.schedule_policy import PrefillAdder


class TestRemTotalTokenOffsetBug(unittest.TestCase):
    """Test to expose the bug at schedule_policy.py:700"""
    
    def test_preemption_with_different_token_counts(self):
        """
        Test that demonstrates the bug when preempting requests with 
        different max_new_tokens values.
        """
        # Setup: Create mock objects
        tree_cache = Mock()
        tree_cache.evictable_size = Mock(return_value=10000)
        
        token_allocator = Mock()
        token_allocator.available_size = Mock(return_value=10000)
        
        # Create a running request with 1000 max_new_tokens
        running_req = Mock()
        running_req.sampling_params = Mock()
        running_req.sampling_params.max_new_tokens = 1000
        running_req.output_ids = []  # No tokens generated yet
        
        # Create running batch with the running request
        running_batch = Mock()
        running_batch.reqs = [running_req]
        
        # Initialize PrefillAdder
        adder = PrefillAdder(
            page_size=16,
            tree_cache=tree_cache,
            token_to_kv_pool_allocator=token_allocator,
            running_batch=running_batch,
            new_token_ratio=1.0,
            rem_input_tokens=10000,
            rem_chunk_tokens=10000,
            mixed_with_decode_tokens=0,
            priority_scheduling_preemption_threshold=10,
        )
        
        # Initial offset should include the running request's tokens
        initial_offset = adder.rem_total_token_offset
        print(f"Initial rem_total_token_offset: {initial_offset}")
        assert initial_offset == 1000, f"Expected 1000, got {initial_offset}"
        
        # Now create a new incoming request with 5000 max_new_tokens
        new_req = Mock()
        new_req.sampling_params = Mock()
        new_req.sampling_params.max_new_tokens = 5000
        new_req.output_ids = []
        
        # Simulate preemption: This is where the bug occurs
        # The code at line 700 incorrectly uses 'req' instead of 'running_req'
        
        # Bug simulation: What the current code does (WRONG)
        wrong_offset_reduction = adder._get_running_request_total_token_offset(new_req)
        print(f"BUG: Code uses new_req and reduces by: {wrong_offset_reduction}")
        
        # Correct behavior: What it should do (RIGHT)
        correct_offset_reduction = adder._get_running_request_total_token_offset(running_req)
        print(f"CORRECT: Should use running_req and reduce by: {correct_offset_reduction}")
        
        # Demonstrate the difference
        assert wrong_offset_reduction == 5000, "Bug calculation should be 5000"
        assert correct_offset_reduction == 1000, "Correct calculation should be 1000"
        
        # Show the impact
        print("\n=== Impact Analysis ===")
        print(f"Wrong calculation reduces offset by: {wrong_offset_reduction}")
        print(f"Correct calculation should reduce by: {correct_offset_reduction}")
        print(f"Difference: {wrong_offset_reduction - correct_offset_reduction} tokens")
        print("\nThis leads to:")
        print("- System thinks it has 4000 MORE tokens available than reality")
        print("- Risk: Over-commitment → OOM (Out of Memory)")
        
        # Assert the bug exists
        assert wrong_offset_reduction != correct_offset_reduction, \
            "BUG CONFIRMED: Wrong variable causes incorrect offset calculation!"


if __name__ == "__main__":
    # Run the test
    test = TestRemTotalTokenOffsetBug()
    test.test_preemption_with_different_token_counts()
    print("\n✅ Bug successfully reproduced!")