"""Task2.9 跨机 KV Transfer：Socket backend 集成测试。"""

from __future__ import annotations

import socket
from contextlib import closing
from typing import Any

from sagellm_kv_cache import KVHandle, KVTransferEngine, SocketTransferBackend


def _free_port() -> int:
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def _build_endpoints(world_size: int) -> dict[int, tuple[str, int]]:
    return {rank: ("127.0.0.1", _free_port()) for rank in range(world_size)}


def test_socket_backend_list_round_trip() -> None:
    endpoints = _build_endpoints(world_size=2)
    backend0 = SocketTransferBackend(rank=0, world_size=2, endpoints=endpoints)
    backend1 = SocketTransferBackend(rank=1, world_size=2, endpoints=endpoints)

    backend0.start()
    backend1.start()
    try:
        sender = KVTransferEngine(backend0)
        receiver = KVTransferEngine(backend1)

        handle = KVHandle.create(num_tokens=16, request_id="req-socket-list")
        recv_buffer: list[Any] = []

        sender.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[1, 2, 3])
        receiver.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=recv_buffer)

        assert recv_buffer == [1, 2, 3]
    finally:
        backend0.close()
        backend1.close()


def test_socket_backend_bytes_compression_round_trip() -> None:
    endpoints = _build_endpoints(world_size=2)
    backend0 = SocketTransferBackend(rank=0, world_size=2, endpoints=endpoints)
    backend1 = SocketTransferBackend(rank=1, world_size=2, endpoints=endpoints)

    backend0.start()
    backend1.start()
    try:
        sender = KVTransferEngine(backend0, enable_compression=True, enable_checksum=True)
        receiver = KVTransferEngine(backend1, enable_compression=True, enable_checksum=False)

        original = bytearray(b"cross-node-payload-" * 256)
        recv_buffer = bytearray(len(original))
        handle = KVHandle.create(num_tokens=64, request_id="req-socket-bytes")

        sender.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=original)
        receiver.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=recv_buffer)

        assert bytes(recv_buffer) == bytes(original)
    finally:
        backend0.close()
        backend1.close()


def test_socket_backend_metrics_observable() -> None:
    endpoints = _build_endpoints(world_size=2)
    backend0 = SocketTransferBackend(rank=0, world_size=2, endpoints=endpoints)
    backend1 = SocketTransferBackend(rank=1, world_size=2, endpoints=endpoints)

    backend0.start()
    backend1.start()
    try:
        sender = KVTransferEngine(backend0, enable_compression=True, max_retries=1)
        receiver = KVTransferEngine(backend1, enable_compression=True)

        handle = KVHandle.create(num_tokens=8, request_id="req-socket-metrics")
        payload = b"A" * 4096
        recv_buffer = bytearray(len(payload))

        tid = sender.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        receiver.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=recv_buffer)

        metrics = sender.get_transfer_metrics(tid)
        assert metrics["bytes"] > 0
        assert metrics["latency_ms"] >= 0.0
        assert metrics["bandwidth_gbps"] >= 0.0
        assert metrics["status"] == "completed"
        assert metrics["compression_ratio"] >= 1.0
    finally:
        backend0.close()
        backend1.close()
