"""Integration tests for IR + KV Cache coordination."""

from __future__ import annotations

import hashlib
from datetime import datetime
from types import SimpleNamespace

import pytest
from sagellm_kv_cache import SchedulerBridge

from sagellm_control.ir.builder import IRBuilder
from sagellm_control.ir.executor import DefaultIRExecutor
from sagellm_control.ir.optimizer import KVReusePass
from sagellm_control.policies import KVAwareIRPolicy
from sagellm_control.types import EngineInfo, EngineState, RequestMetadata, RequestPriority


class _InMemoryEngineClient:
    async def execute(self, request):
        return SimpleNamespace(
            request_id=request.request_id,
            output_text=f"ok-{request.request_id}",
            finish_reason="stop",
        )


def _stable_token_ids(prompt: str) -> list[int]:
    token_ids: list[int] = []
    for word in prompt.split():
        digest = hashlib.sha1(word.encode("utf-8")).digest()
        token_ids.append(int.from_bytes(digest[:2], byteorder="big", signed=False))
    return token_ids


def test_ir_builder_queries_kv_cache_and_adjusts_budget() -> None:
    """Builder should query KV availability and keep allocation within budget."""
    bridge = SchedulerBridge(max_tokens=32, block_size=4, enable_prefix_cache=True)
    builder = IRBuilder(scheduler_bridge=bridge)

    meta = RequestMetadata(
        request_id="req-kv-budget",
        trace_id="trace-kv-budget",
        model_id="Qwen2-7B",
        prompt="one two three four five six seven eight nine ten",
        max_tokens=40,
        prompt_tokens=10,
        priority=RequestPriority.NORMAL,
    )

    ir = builder.build_from_request(meta)

    prefill = ir.get_node("prefill-req-kv-budget")
    decode = ir.get_node("decode-req-kv-budget")
    assert prefill is not None
    assert decode is not None

    allocated = (
        prefill.resource_allocation.kv_cache_tokens + decode.resource_allocation.kv_cache_tokens
    )
    assert allocated <= 32
    assert ir.optimization_hints["kv_bridge_enabled"] is True
    assert ir.optimization_hints["kv_budget_exceeded"] is True


def test_ir_builder_and_optimizer_use_prefix_cache() -> None:
    """Builder and KVReusePass should mark reusable prefix from cache."""
    bridge = SchedulerBridge(max_tokens=256, block_size=2, enable_prefix_cache=True)
    prompt = "prefix shared tokens alpha beta"

    token_ids = _stable_token_ids(prompt)
    block_hashes = bridge.prefix_cache.compute_block_hashes(token_ids)
    cached_blocks = [{"block_id": idx} for idx in range(len(block_hashes))]
    bridge.prefix_cache.insert(block_hashes, cached_blocks)

    builder = IRBuilder(scheduler_bridge=bridge)
    ir = builder.build_from_request(
        RequestMetadata(
            request_id="req-prefix",
            trace_id="trace-prefix",
            model_id="Qwen2-7B",
            prompt=prompt,
            max_tokens=8,
            prompt_tokens=len(token_ids),
            priority=RequestPriority.NORMAL,
        )
    )

    prefill = ir.get_node("prefill-req-prefix")
    assert prefill is not None
    assert prefill.resource_allocation.reuse_kv_prefix is True
    assert prefill.resource_allocation.kv_prefix_length > 0

    pass_ = KVReusePass(min_reuse_length=1, kv_bridge=bridge)
    optimized = pass_.optimize(ir)
    assert optimized.optimization_hints["kv_cache_hits"] >= 1
    assert optimized.optimization_hints["kv_cache_hit_rate"] > 0


@pytest.mark.asyncio
async def test_ir_executor_kv_allocation_and_miss_fallback() -> None:
    """Executor should coordinate KV allocation and still fallback on miss."""
    bridge = SchedulerBridge(max_tokens=24, block_size=4, enable_prefix_cache=True)
    builder = IRBuilder(scheduler_bridge=bridge)
    executor = DefaultIRExecutor(
        engine_client=_InMemoryEngineClient(),
        scheduler_bridge=bridge,
    )

    ok_ir = builder.build_from_request(
        RequestMetadata(
            request_id="req-kv-ok",
            trace_id="trace-kv-ok",
            model_id="Qwen2-7B",
            prompt="short prompt",
            max_tokens=6,
            prompt_tokens=2,
            priority=RequestPriority.NORMAL,
        )
    )
    ok_results = await executor.execute(ok_ir)
    assert "cmd-req-kv-ok" in ok_results
    assert ok_results["cmd-req-kv-ok"]["output_text"] == "ok-req-kv-ok"

    miss_ir = builder.build_from_request(
        RequestMetadata(
            request_id="req-kv-miss",
            trace_id="trace-kv-miss",
            model_id="Qwen2-7B",
            prompt="very long prompt that exceeds kv allocation budget for fallback behavior",
            max_tokens=100,
            prompt_tokens=12,
            priority=RequestPriority.NORMAL,
        )
    )
    miss_results = await executor.execute(miss_ir)
    assert "cmd-req-kv-miss" in miss_results
    assert miss_results["cmd-req-kv-miss"]["output_text"] == "ok-req-kv-miss"


def test_kv_aware_ir_policy_prefers_reusable_ir() -> None:
    """KVAwareIRPolicy should prioritize IRs with reusable prefix and fitting budget."""
    policy = KVAwareIRPolicy(max_kv_tokens_per_engine=128)

    req_a = RequestMetadata(
        request_id="req-a",
        trace_id="trace-a",
        model_id="Qwen2-7B",
        prompt="aaa",
        max_tokens=16,
        prompt_tokens=1,
        priority=RequestPriority.NORMAL,
        arrival_time=datetime.now(),
    )
    req_b = RequestMetadata(
        request_id="req-b",
        trace_id="trace-b",
        model_id="Qwen2-7B",
        prompt="bbb",
        max_tokens=16,
        prompt_tokens=1,
        priority=RequestPriority.NORMAL,
        arrival_time=datetime.now(),
    )

    ir_by_request_id = {
        "req-a": SimpleNamespace(
            optimization_hints={
                "prefix_cache_reusable_tokens": 8,
                "kv_required_tokens": 20,
                "kv_available_tokens": 40,
            }
        ),
        "req-b": SimpleNamespace(
            optimization_hints={
                "prefix_cache_reusable_tokens": 0,
                "kv_required_tokens": 60,
                "kv_available_tokens": 40,
            }
        ),
    }

    ordered = policy.reorder_ir_requests([req_b, req_a], ir_by_request_id)
    assert ordered[0].request_id == "req-a"


def test_kv_aware_ir_policy_schedule_and_utilization_metrics() -> None:
    """KVAwareIRPolicy should schedule with KV budget awareness and track utilization."""
    policy = KVAwareIRPolicy(max_kv_tokens_per_engine=64)

    candidate = EngineInfo(
        engine_id="engine-1",
        model_id="Qwen2-7B",
        host="127.0.0.1",
        port=8000,
        state=EngineState.READY,
    )

    request = RequestMetadata(
        request_id="req-util",
        trace_id="trace-util",
        model_id="Qwen2-7B",
        prompt="budget aware request",
        max_tokens=8,
        prompt_tokens=4,
        priority=RequestPriority.NORMAL,
    )

    decision = policy.schedule(request, [candidate], [])
    assert decision.engine_id == "engine-1"
    metrics = policy.get_kv_utilization_metrics()
    assert "engine-1" in metrics
    assert metrics["engine-1"] > 0.0
