"""测试 PD 分离功能

验证 CPUEngine 的 prefill() 和 decode() 方法是否正常工作。
"""

from __future__ import annotations

import asyncio
import logging

import pytest

from sagellm_protocol.types import Request
from sagellm_core import PDSeparatedExecutor, DistributedRuntime, create_backend, create_engine
from sagellm_core.config import EngineConfig

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@pytest.mark.asyncio
async def test_cpu_engine_prefill():
    """测试 CPUEngine 的 prefill 方法"""
    # 创建引擎 - 使用轻量级模型（约 50MB，快速下载）
    backend = create_backend({"kind": "cpu"})
    config = EngineConfig(
        kind="cpu",
        model="sshleifer/tiny-gpt2",  # 轻量级测试模型
        device="cpu",
    )
    engine = create_engine(config, backend)

    try:
        await engine.start()

        # 创建请求
        request = Request(
            request_id="req-prefill-001",
            trace_id="trace-001",
            model="sshleifer/tiny-gpt2",
            prompt="Hello, how are you?",
            max_tokens=10,
            stream=False,
        )

        # 执行 Prefill
        result = await engine.prefill(request)

        # 验证结果
        assert "kv_handle" in result
        assert "num_tokens" in result
        assert result["num_tokens"] > 0
        assert "first_token_id" in result

        logger.info(f"✓ Prefill completed: {result['num_tokens']} tokens processed")

    finally:
        await engine.stop()


@pytest.mark.asyncio
async def test_cpu_engine_decode():
    """测试 CPUEngine 的 decode 方法"""
    backend = create_backend({"kind": "cpu"})
    config = EngineConfig(
        kind="cpu",
        model="sshleifer/tiny-gpt2",
        device="cpu",
    )
    engine = create_engine(config, backend)

    try:
        await engine.start()

        # 先执行 Prefill
        request = Request(
            request_id="req-decode-001",
            trace_id="trace-001",
            model="sshleifer/tiny-gpt2",
            prompt="Hello!",
            max_tokens=10,
            stream=False,
        )

        prefill_result = await engine.prefill(request)
        kv_handle = prefill_result["kv_handle"]

        # 执行 Decode
        decode_result = await engine.decode(request, kv_handle=kv_handle, max_new_tokens=10)

        # 验证结果
        assert "output_tokens" in decode_result
        assert "output_text" in decode_result
        assert "finish_reason" in decode_result
        assert len(decode_result["output_tokens"]) > 0

        logger.info(
            f"✓ Decode completed: {len(decode_result['output_tokens'])} tokens generated, "
            f"text='{decode_result['output_text']}'"
        )

    finally:
        await engine.stop()


@pytest.mark.asyncio
async def test_pd_executor_hybrid():
    """测试 PDSeparatedExecutor 的 Hybrid 模式"""
    # 初始化 Runtime
    runtime = DistributedRuntime()
    await runtime.initialize()

    # 创建引擎
    backend = create_backend({"kind": "cpu"})
    config = EngineConfig(
        kind="cpu",
        model="sshleifer/tiny-gpt2",
        device="cpu",
    )
    engine = create_engine(config, backend)

    try:
        await engine.start()

        # 创建 PD Executor
        executor = PDSeparatedExecutor(engine=engine, runtime=runtime)

        # 创建 Hybrid 请求
        request = Request(
            request_id="req-hybrid-001",
            trace_id="trace-001",
            model="sshleifer/tiny-gpt2",
            prompt="Hello!",
            max_tokens=10,
            stream=False,
            phase="decode",  # hybrid mode
        )

        # 执行
        response = await executor.execute(request)

        # 验证结果
        assert response.request_id == "req-hybrid-001"
        assert response.metrics.prefill_ms > 0
        assert response.metrics.decode_ms > 0
        assert len(response.output_text) > 0

        logger.info("✓ Hybrid execution completed:")
        logger.info(f"  - Prefill: {response.metrics.prefill_ms:.2f} ms")
        logger.info(f"  - Decode: {response.metrics.decode_ms:.2f} ms")
        logger.info(f"  - Output: '{response.output_text}'")

    finally:
        await engine.stop()
        await runtime.shutdown()


@pytest.mark.asyncio
async def test_pd_executor_prefill_only():
    """测试 PDSeparatedExecutor 的 Prefill-Only 模式"""
    runtime = DistributedRuntime()
    await runtime.initialize()

    backend = create_backend({"kind": "cpu"})
    config = EngineConfig(
        kind="cpu",
        model="sshleifer/tiny-gpt2",
        device="cpu",
    )
    engine = create_engine(config, backend)

    try:
        await engine.start()
        executor = PDSeparatedExecutor(engine=engine, runtime=runtime)

        # Prefill-Only 请求
        request = Request(
            request_id="req-prefill-only-001",
            trace_id="trace-001",
            model="sshleifer/tiny-gpt2",
            prompt="Hello!",
            max_tokens=1,  # 最小值（但实际只做 prefill）
            stream=False,
            phase="prefill",
        )

        response = await executor.execute(request)

        # 验证：只有 Prefill，没有 Decode（或 Decode 很少）
        assert response.metrics.prefill_ms > 0
        # decode_ms 可能为 0（prefill-only）或很小（生成了 1 token）

        logger.info("✓ Prefill-Only execution completed:")
        logger.info(f"  - Prefill: {response.metrics.prefill_ms:.2f} ms")

    finally:
        await engine.stop()
        await runtime.shutdown()


if __name__ == "__main__":
    # 运行测试
    asyncio.run(test_cpu_engine_prefill())
    asyncio.run(test_cpu_engine_decode())
    asyncio.run(test_pd_executor_hybrid())
    asyncio.run(test_pd_executor_prefill_only())
    print("\n✅ All PD separation tests passed!")
