"""
Simple end-to-end runtime execution test to verify context cleanup fixes.

Run with: python tests/test_runtime_execution_e2e.py
"""

import asyncio
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from control_plane_api.worker.runtimes.claude_code.litellm_proxy import (
    set_execution_context,
    clear_execution_context,
    _context_store
)


async def test_execution_with_cleanup():
    """Test that mimics a real execution lifecycle with context cleanup."""
    print("\n🧪 Testing execution lifecycle with delayed cleanup...")

    # Simulate execution starting
    execution_id = "test-exec-123"
    print(f"  ✓ Starting execution: {execution_id}")

    set_execution_context(execution_id, {
        "user_id": "user-123",
        "session_id": "session-456",
        "organization_id": "org-789",
        "agent_id": "agent-abc"
    })

    # Verify context is set
    context = _context_store.get_context(execution_id)
    assert context is not None, "Context should be set"
    print(f"  ✓ Context set successfully")

    # Simulate some execution work
    await asyncio.sleep(0.1)

    # Simulate execution completing - schedule delayed cleanup
    print(f"  ✓ Execution completed, scheduling cleanup with 2s delay...")
    clear_execution_context(
        execution_id,
        immediate=False,
        delay_seconds=2.0
    )

    # Context should still be available immediately after scheduling cleanup
    context = _context_store.get_context(execution_id)
    assert context is not None, "Context should still be available immediately after cleanup scheduled"
    print(f"  ✓ Context still available immediately after scheduling cleanup")

    # Simulate SDK request that happens 0.5s after execution completes
    await asyncio.sleep(0.5)
    context = _context_store.get_context(execution_id)
    assert context is not None, "Context should still be available during in-flight SDK request"
    print(f"  ✓ Context available for in-flight SDK request (0.5s after completion)")

    # Wait for cleanup to execute
    print(f"  ⏳ Waiting for cleanup to execute (2s delay)...")
    await asyncio.sleep(1.7)  # Total 2.2s - cleanup should have happened

    # Context should be cleared now
    context = _context_store.get_context(execution_id)
    assert context is None, "Context should be cleared after delay"
    print(f"  ✓ Context cleared after delay")

    print("✅ All checks passed!\n")


async def test_concurrent_executions():
    """Test multiple concurrent executions don't interfere."""
    print("\n🧪 Testing concurrent executions...")

    async def execution_lifecycle(exec_id: str):
        """Single execution lifecycle."""
        set_execution_context(exec_id, {
            "user_id": f"user-{exec_id}",
            "organization_id": "org-test"
        })

        # Do some work
        await asyncio.sleep(0.1)

        # Schedule cleanup
        clear_execution_context(exec_id, immediate=False, delay_seconds=0.5)

        # Verify context still available
        ctx = _context_store.get_context(exec_id)
        assert ctx is not None, f"Context for {exec_id} should be available"

        # Wait for cleanup
        await asyncio.sleep(0.6)

        # Verify cleaned up
        ctx = _context_store.get_context(exec_id)
        assert ctx is None, f"Context for {exec_id} should be cleared"

    # Run 5 concurrent executions
    tasks = [execution_lifecycle(f"exec-{i}") for i in range(5)]
    await asyncio.gather(*tasks)

    print("  ✓ 5 concurrent executions completed without interference")
    print("✅ Concurrent execution test passed!\n")


async def test_ttl_cleanup():
    """Test that TTL cleanup prevents memory leaks."""
    print("\n🧪 Testing TTL cleanup...")

    from control_plane_api.worker.runtimes.claude_code.litellm_proxy import ExecutionContextStore

    # Create store with 1-second TTL
    store = ExecutionContextStore(ttl_seconds=1)

    # Add some contexts
    for i in range(10):
        store.set_context(f"exec-{i}", {"user_id": f"user-{i}"})

    print(f"  ✓ Added 10 contexts")

    # All should be available
    for i in range(10):
        assert store.get_context(f"exec-{i}") is not None

    print(f"  ✓ All contexts accessible")

    # Wait for TTL to expire
    await asyncio.sleep(1.5)

    # Add new context (triggers cleanup)
    store.set_context("exec-new", {"user_id": "user-new"})

    print(f"  ✓ Waited for TTL expiry and added new context")

    # Old contexts should be expired
    expired_count = sum(1 for i in range(10) if store.get_context(f"exec-{i}") is None)
    assert expired_count == 10, f"All 10 old contexts should be expired, but {expired_count} were expired"

    # New context should be available
    assert store.get_context("exec-new") is not None

    print(f"  ✓ All old contexts expired, new context available")
    print("✅ TTL cleanup test passed!\n")


async def main():
    """Run all tests."""
    print("\n" + "="*60)
    print("  Worker Context Cleanup End-to-End Tests")
    print("="*60)

    try:
        await test_execution_with_cleanup()
        await test_concurrent_executions()
        await test_ttl_cleanup()

        print("\n" + "="*60)
        print("  ✅ ALL TESTS PASSED!")
        print("="*60 + "\n")
        return 0
    except AssertionError as e:
        print(f"\n❌ Test failed: {e}\n")
        return 1
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}\n")
        import traceback
        traceback.print_exc()
        return 1


if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)
