"""KV Transfer tests for Task2.8/2.9.

Task2.8（baseline）使用真实可执行用例验证 metadata + payload 的最小闭环。
Task2.9（cross-node advanced）仍保留占位测试。
"""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock

import pytest
from sagellm_protocol.kv.enums import TransferDirection, TransferStatus

from sagellm_kv_cache import DType, KVHandle, KVTransferEngine, KVTransferFailedError


class _InMemoryBackend:
    """Minimal in-memory backend for cross-node transfer tests."""

    def __init__(self, rank: int, mailbox: dict[tuple[int, int], list[Any]]) -> None:
        self._rank = rank
        self._mailbox = mailbox

    def get_rank(self) -> int:
        return self._rank

    def send(self, tensor: Any, dst_rank: int) -> None:
        key = (self._rank, dst_rank)
        if key not in self._mailbox:
            self._mailbox[key] = []
        self._mailbox[key].append(tensor)

    def recv(self, tensor: list[Any], src_rank: int) -> None:
        key = (src_rank, self._rank)
        queue = self._mailbox.get(key, [])
        if not queue:
            raise RuntimeError(f"no payload from {src_rank} to {self._rank}")
        tensor.clear()
        tensor.extend(queue.pop(0))


class _FailingCongestionBackend(_InMemoryBackend):
    """Backend that simulates congestion on send."""

    def send(self, tensor: Any, dst_rank: int) -> None:
        raise RuntimeError(f"congestion: dst_rank={dst_rank}")


class TestKVTransferProtocol:
    """Task2.8 - KV Transfer Protocol baseline tests."""

    def test_kv_transfer_metadata_schema(self, cpu_config: dict[str, Any]) -> None:
        """Verify metadata schema is populated after transfer."""
        backend = MagicMock()
        backend.get_rank.return_value = 0
        engine = KVTransferEngine(backend)

        handle = KVHandle.create(
            num_tokens=64,
            dtype=DType.FP16,
            device="cpu",
            request_id="req-metadata",
        )
        tensor = MagicMock()

        transfer_id = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=tensor)
        metadata = engine.get_transfer_metadata(transfer_id)

        assert metadata.transfer_id == transfer_id
        assert metadata.request_id == "req-metadata"
        assert metadata.token_start == 0
        assert metadata.token_end == 64
        assert metadata.total_bytes > 0
        assert metadata.direction == TransferDirection.SEND
        assert metadata.status == TransferStatus.COMPLETED

    def test_send_kv_payload_path(self, cpu_config: dict[str, Any]) -> None:
        """Verify sender path invokes backend send with payload."""
        backend = MagicMock()
        backend.get_rank.return_value = 0
        engine = KVTransferEngine(backend)

        handle = KVHandle.create(num_tokens=32, request_id="req-send")
        tensor = MagicMock()

        transfer_id = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=tensor)

        backend.send.assert_called_once_with(tensor, 1)
        assert engine.get_transfer_status(transfer_id) == TransferStatus.COMPLETED

    def test_recv_kv_payload_path(self, cpu_config: dict[str, Any]) -> None:
        """Verify receiver path invokes backend recv into payload buffer."""
        backend = MagicMock()
        backend.get_rank.return_value = 1
        engine = KVTransferEngine(backend)

        handle = KVHandle.create(num_tokens=32, request_id="req-recv")
        recv_buffer = MagicMock()

        transfer_id = engine.transfer_kv_handle(
            handle,
            src_rank=0,
            dst_rank=1,
            tensor=recv_buffer,
        )

        backend.recv.assert_called_once_with(recv_buffer, 0)
        assert engine.get_transfer_status(transfer_id) == TransferStatus.COMPLETED

    def test_kv_shard_transfer_baseline(self, cpu_config: dict[str, Any]) -> None:
        """Verify baseline shard-style transfer using two handle slices."""
        backend = MagicMock()
        backend.get_rank.return_value = 0
        engine = KVTransferEngine(backend)

        shard1 = KVHandle.create(num_tokens=16, request_id="req-shard")
        shard1.token_start = 0
        shard1.token_end = 16

        shard2 = KVHandle.create(num_tokens=16, request_id="req-shard")
        shard2.token_start = 16
        shard2.token_end = 32

        tensor1 = MagicMock()
        tensor2 = MagicMock()

        t1 = engine.transfer_kv_handle(shard1, src_rank=0, dst_rank=1, tensor=tensor1)
        t2 = engine.transfer_kv_handle(shard2, src_rank=0, dst_rank=1, tensor=tensor2)

        m1 = engine.get_transfer_metadata(t1)
        m2 = engine.get_transfer_metadata(t2)

        assert (m1.token_start, m1.token_end) == (0, 16)
        assert (m2.token_start, m2.token_end) == (16, 32)
        assert engine.get_completed_count() == 2

    def test_compression_hook_baseline_disabled(self, cpu_config: dict[str, Any]) -> None:
        """Verify baseline keeps compression disabled by default."""
        backend = MagicMock()
        backend.get_rank.return_value = 0
        engine = KVTransferEngine(backend)

        handle = KVHandle.create(num_tokens=8, request_id="req-compress")
        tensor = MagicMock()

        transfer_id = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=tensor)
        metadata = engine.get_transfer_metadata(transfer_id)

        assert metadata.is_compressed is False
        assert metadata.original_size == 0


class TestCrossNodeKVTransfer:
    """Task2.9 - Cross-node KV Transfer baseline tests."""

    def test_cross_node_transfer_baseline(self, cpu_config: dict[str, Any]) -> None:
        """Verify end-to-end transfer across two in-memory nodes."""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        sender_engine = KVTransferEngine(_InMemoryBackend(rank=0, mailbox=mailbox))
        receiver_engine = KVTransferEngine(_InMemoryBackend(rank=1, mailbox=mailbox))

        handle = KVHandle.create(num_tokens=16, request_id="req-cross-node")
        payload = [1, 2, 3, 4]
        recv_buffer: list[int] = []

        sender_tid = sender_engine.transfer_kv_handle(
            handle, src_rank=0, dst_rank=1, tensor=payload
        )
        receiver_tid = receiver_engine.transfer_kv_handle(
            handle,
            src_rank=0,
            dst_rank=1,
            tensor=recv_buffer,
        )

        assert recv_buffer == payload
        assert sender_engine.get_transfer_status(sender_tid) == TransferStatus.COMPLETED
        assert receiver_engine.get_transfer_status(receiver_tid) == TransferStatus.COMPLETED

    def test_rdma_backend_adapter_contract_baseline(self, cpu_config: dict[str, Any]) -> None:
        """Verify an RDMA-like backend that honors send/recv contract works with engine."""
        mailbox: dict[tuple[int, int], list[Any]] = {}

        class FakeRDMABackend(_InMemoryBackend):
            pass

        sender_engine = KVTransferEngine(FakeRDMABackend(rank=2, mailbox=mailbox))
        receiver_engine = KVTransferEngine(FakeRDMABackend(rank=3, mailbox=mailbox))

        handle = KVHandle.create(num_tokens=8, request_id="req-rdma")
        payload = [7, 8, 9]
        recv_buffer: list[int] = []

        sender_engine.transfer_kv_handle(handle, src_rank=2, dst_rank=3, tensor=payload)
        receiver_engine.transfer_kv_handle(handle, src_rank=2, dst_rank=3, tensor=recv_buffer)

        assert recv_buffer == payload

    def test_compression_integration_baseline_disabled(self, cpu_config: dict[str, Any]) -> None:
        """Verify cross-node baseline keeps compression fields at default values."""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        sender_engine = KVTransferEngine(_InMemoryBackend(rank=4, mailbox=mailbox))

        handle = KVHandle.create(num_tokens=10, request_id="req-compression")
        payload = [42]

        transfer_id = sender_engine.transfer_kv_handle(
            handle, src_rank=4, dst_rank=5, tensor=payload
        )
        metadata = sender_engine.get_transfer_metadata(transfer_id)

        assert metadata.is_compressed is False
        assert metadata.original_size == 0
        assert metadata.checksum is None

    def test_congestion_control_failure_path(self, cpu_config: dict[str, Any]) -> None:
        """Verify send-path congestion surfaces a structured transfer failure."""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(_FailingCongestionBackend(rank=6, mailbox=mailbox))

        handle = KVHandle.create(num_tokens=12, request_id="req-congestion")
        payload = [99]

        with pytest.raises(KVTransferFailedError, match="congestion"):
            engine.transfer_kv_handle(handle, src_rank=6, dst_rank=7, tensor=payload)
