"""End-to-end test for KVCacheManager integration into LLMEngine.

Tests that KVCacheManager is properly integrated throughout the inference pipeline:
LLMEngine → Executor → Worker → ModelRunner

Issue: https://github.com/intellistream/sagellm-core/issues/33
"""

from __future__ import annotations

import pytest

from sagellm_core import LLMEngine, LLMEngineConfig


@pytest.mark.asyncio
async def test_kv_cache_manager_initialized():
    """Test that KVCacheManager is initialized by LLMEngine."""
    config = LLMEngineConfig(
        model_path="sshleifer/tiny-gpt2",
        backend_type="cpu",
        max_new_tokens=10,
    )
    engine = LLMEngine(config)

    # Before start, KV cache should be None
    assert engine.get_kv_cache() is None

    # Start engine
    await engine.start()

    # After start, KV cache should be initialized
    kv_cache = engine.get_kv_cache()
    assert kv_cache is not None
    assert hasattr(kv_cache, "allocate_slots")
    assert hasattr(kv_cache, "free")
    assert hasattr(kv_cache, "can_alloc")

    # Cleanup
    await engine.stop()


@pytest.mark.asyncio
async def test_kv_cache_manager_used_in_inference():
    """Test that KVCacheManager is actually used during inference."""
    config = LLMEngineConfig(
        model_path="sshleifer/tiny-gpt2",
        backend_type="cpu",
        max_new_tokens=10,
    )
    engine = LLMEngine(config)
    await engine.start()

    kv_cache = engine.get_kv_cache()
    assert kv_cache is not None

    # Get initial stats
    initial_stats = kv_cache.get_stats()
    initial_num_requests = initial_stats.get("num_requests", 0)

    # Run inference (should allocate and free KV cache)
    response = await engine.generate(
        prompt="Hello, how are",
        max_tokens=5,
    )

    assert response is not None
    assert response.output_text is not None

    # After inference, KV cache should have been used
    # Note: Since we free KV cache after inference, num_requests should be back to 0
    final_stats = kv_cache.get_stats()
    final_num_requests = final_stats.get("num_requests", 0)

    # The request should have been allocated and freed
    assert final_num_requests == initial_num_requests

    # Cleanup
    await engine.stop()


@pytest.mark.asyncio
async def test_kv_cache_allocation_capacity():
    """Test that KVCacheManager enforces capacity limits."""
    config = LLMEngineConfig(
        model_path="sshleifer/tiny-gpt2",
        backend_type="cpu",
        max_new_tokens=50,
    )
    engine = LLMEngine(config)
    await engine.start()

    kv_cache = engine.get_kv_cache()
    assert kv_cache is not None

    # Check that we can allocate within capacity
    # Default config: 512 blocks × 128 tokens/block = 65536 tokens
    assert kv_cache.can_alloc(1000)
    assert kv_cache.can_alloc(10000)

    # Should not be able to allocate more than max_tokens
    assert not kv_cache.can_alloc(100000)

    # Cleanup
    await engine.stop()


@pytest.mark.asyncio
async def test_multiple_requests_kv_cache():
    """Test that KVCacheManager handles multiple sequential requests."""
    config = LLMEngineConfig(
        model_path="sshleifer/tiny-gpt2",
        backend_type="cpu",
        max_new_tokens=5,
    )
    engine = LLMEngine(config)
    await engine.start()

    kv_cache = engine.get_kv_cache()
    assert kv_cache is not None

    # Run multiple requests sequentially
    prompts = [
        "Hello world",
        "How are you",
        "What is AI",
    ]

    for prompt in prompts:
        response = await engine.generate(prompt=prompt, max_tokens=5)
        assert response is not None
        assert response.output_text is not None

    # All requests should have been cleaned up
    stats = kv_cache.get_stats()
    assert stats.get("num_requests", 0) == 0

    # Cleanup
    await engine.stop()


if __name__ == "__main__":
    import asyncio

    async def main():
        """Run tests manually."""
        print("Test 1: KV cache manager initialized")
        await test_kv_cache_manager_initialized()
        print("✓ PASSED\n")

        print("Test 2: KV cache manager used in inference")
        await test_kv_cache_manager_used_in_inference()
        print("✓ PASSED\n")

        print("Test 3: KV cache allocation capacity")
        await test_kv_cache_allocation_capacity()
        print("✓ PASSED\n")

        print("Test 4: Multiple requests KV cache")
        await test_multiple_requests_kv_cache()
        print("✓ PASSED\n")

        print("All tests passed! ✓")

    asyncio.run(main())
