"""Task2.9 Cross-node KV Transfer 完整验收测试套件

覆盖：
1. 压缩/解压链路（bytes payload 的 zlib 压缩，metadata 字段验证）
2. CRC32 校验和（send 端计算，recv 端记录）
3. 拥塞控制/重试（指数退避，retry_count 可观测）
4. 端到端稳定性基准（多次传输，延迟收敛，指标完整）

运行：
    pytest tests/test_kv_transfer_crossnode_advanced.py -v
"""

from __future__ import annotations

import zlib
from typing import Any

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

from sagellm_kv_cache import KVHandle, KVTransferEngine, KVTransferFailedError

# ---------------------------------------------------------------------------
# 测试用 backend：支持 bytes/bytearray/list 载荷的内存邮箱
# ---------------------------------------------------------------------------


class _AdvancedMailboxBackend:
    """支持 bytes/bytearray/list 的内存邮箱 backend（用于压缩链路测试）。"""

    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] = []
        # 对 bytearray/bytes 做深拷贝，避免引用污染
        if isinstance(tensor, (bytes, bytearray)):
            self._mailbox[key].append(bytes(tensor))
        else:
            self._mailbox[key].append(tensor)

    def recv(self, tensor: 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}")
        payload = queue.pop(0)
        if isinstance(tensor, bytearray):
            tensor[:] = payload
        elif isinstance(tensor, list):
            tensor.clear()
            tensor.extend(payload if isinstance(payload, list) else [payload])
        else:
            raise TypeError(f"Unsupported recv buffer type: {type(tensor)}")


class _TransientFailBackend(_AdvancedMailboxBackend):
    """前 fail_count 次 send 抛出拥塞异常，之后正常。"""

    def __init__(
        self,
        rank: int,
        mailbox: dict[tuple[int, int], list[Any]],
        fail_count: int = 2,
    ) -> None:
        super().__init__(rank, mailbox)
        self._fail_count = fail_count
        self._send_calls = 0

    def send(self, tensor: Any, dst_rank: int) -> None:
        self._send_calls += 1
        if self._send_calls <= self._fail_count:
            raise RuntimeError(f"simulated congestion (attempt {self._send_calls})")
        super().send(tensor, dst_rank)


# ---------------------------------------------------------------------------
# 1. 压缩/解压链路
# ---------------------------------------------------------------------------


class TestCompressionChain:
    """压缩链路：bytes 载荷压缩、元数据字段、解压还原。"""

    def test_compression_metadata_fields(self) -> None:
        """开启压缩后，is_compressed/original_size/compression_type/total_bytes 应正确设置。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
            compression_level=6,
        )

        payload = b"x" * 1024  # 1 KiB 可压缩数据
        handle = KVHandle.create(num_tokens=16, request_id="req-compress-meta")

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

        assert meta.is_compressed is True
        assert meta.original_size == 1024
        assert meta.compression_type == CompressionType.ZSTD
        assert meta.total_bytes < meta.original_size, "压缩后应小于原始大小"
        assert meta.status == TransferStatus.COMPLETED

    def test_compressed_payload_is_valid_zlib(self) -> None:
        """实际写入邮箱的载荷应是有效的 zlib 压缩数据。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
        )
        payload = b"hello " * 200  # repeating bytes → high compression ratio

        handle = KVHandle.create(num_tokens=8, request_id="req-compress-valid")
        engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)

        # 邮箱中应存放压缩后的数据
        stored = mailbox[(0, 1)][0]
        # 应能被 zlib 解压
        decompressed = zlib.decompress(stored)
        assert decompressed == payload

    def test_compression_ratio_in_metrics(self) -> None:
        """get_transfer_metrics() 中的 compression_ratio 应 > 1 for compressible data."""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
        )
        payload = b"A" * 4096  # highly compressible

        handle = KVHandle.create(num_tokens=32, request_id="req-ratio")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)

        metrics = engine.get_transfer_metrics(tid)
        assert metrics["compression_ratio"] > 1.0
        assert metrics["status"] == "completed"

    def test_decompression_round_trip(self) -> None:
        """send+recv 端对 bytearray 载荷的完整压缩/解压还原。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        original = bytearray(b"KV data: " * 300)
        recv_buffer = bytearray(len(original))

        sender = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
        )
        receiver = KVTransferEngine(
            _AdvancedMailboxBackend(rank=1, mailbox=mailbox),
            enable_compression=True,
        )

        handle = KVHandle.create(num_tokens=16, request_id="req-roundtrip")
        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), "解压后数据应与原始数据一致"

    def test_no_compression_when_disabled(self) -> None:
        """未开启压缩时，is_compressed 应为 False，original_size==0。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=False,
        )
        payload = b"plain " * 100
        handle = KVHandle.create(num_tokens=8, request_id="req-no-compress")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

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


# ---------------------------------------------------------------------------
# 2. CRC32 校验和验证
# ---------------------------------------------------------------------------


class TestChecksumVerification:
    """CRC32 校验和：send 端设置 checksum，recv 端记录并可在 metrics 中观测。"""

    def test_checksum_set_on_send(self) -> None:
        """send 端开启 enable_checksum 时 metadata.checksum 应为 8 位十六进制串。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_checksum=True,
        )
        payload = b"checksum test data"
        handle = KVHandle.create(num_tokens=4, request_id="req-crc")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

        assert meta.checksum is not None
        assert len(meta.checksum) == 8
        assert all(c in "0123456789abcdef" for c in meta.checksum)

    def test_checksum_matches_crc32(self) -> None:
        """send 端的 checksum 应等于 payload 的 CRC32。"""
        import binascii

        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_checksum=True,
        )
        payload = b"deterministic payload"
        expected_crc = f"{binascii.crc32(payload) & 0xFFFFFFFF:08x}"

        handle = KVHandle.create(num_tokens=4, request_id="req-crc-check")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

        assert meta.checksum == expected_crc

    def test_checksum_in_metrics(self) -> None:
        """get_transfer_metrics() 的 'checksum' 字段非空当 enable_checksum=True。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_checksum=True,
        )
        payload = b"metrics checksum"
        handle = KVHandle.create(num_tokens=4, request_id="req-crc-metric")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        metrics = engine.get_transfer_metrics(tid)

        assert metrics["checksum"] != "", "checksum 字段不应为空"

    def test_no_checksum_when_disabled(self) -> None:
        """未开启 checksum 时 metadata.checksum 应为 None。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_checksum=False,
        )
        payload = b"no checksum"
        handle = KVHandle.create(num_tokens=4, request_id="req-no-crc")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

        assert meta.checksum is None

    def test_compression_and_checksum_combined(self) -> None:
        """压缩 + 校验和同时开启时，meta 字段全部正确填充。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
            enable_checksum=True,
        )
        payload = b"combined test " * 100
        handle = KVHandle.create(num_tokens=8, request_id="req-combined")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

        assert meta.is_compressed is True
        assert meta.original_size == len(payload)
        assert meta.checksum is not None and len(meta.checksum) == 8


# ---------------------------------------------------------------------------
# 3. 拥塞控制：重试与指数退避
# ---------------------------------------------------------------------------


class TestCongestionControlRetry:
    """拥塞控制：transient failure 时重试，retry_count 可观测，超限时抛出异常。"""

    def test_retry_succeeds_after_transient_failures(self) -> None:
        """fail_count=2, max_retries=3 时，第3次 send 成功，transfer 完成。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _TransientFailBackend(rank=0, mailbox=mailbox, fail_count=2),
            max_retries=3,
            retry_backoff_ms=1.0,  # 1 ms，避免测试太慢
        )
        handle = KVHandle.create(num_tokens=8, request_id="req-retry-ok")
        payload = [1, 2, 3]

        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
        meta = engine.get_transfer_metadata(tid)

        assert meta.status == TransferStatus.COMPLETED
        assert engine.get_retry_count(tid) == 2  # 2 次失败 → 2 次重试

    def test_retry_count_in_metrics(self) -> None:
        """get_transfer_metrics()['retries'] 应等于 get_retry_count()。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _TransientFailBackend(rank=0, mailbox=mailbox, fail_count=1),
            max_retries=2,
            retry_backoff_ms=1.0,
        )
        handle = KVHandle.create(num_tokens=4, request_id="req-retry-metric")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[42])
        metrics = engine.get_transfer_metrics(tid)

        assert metrics["retries"] == engine.get_retry_count(tid)
        assert int(metrics["retries"]) == 1

    def test_no_retry_without_failures(self) -> None:
        """无失败时，retry_count 应为 0。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            max_retries=3,
            retry_backoff_ms=1.0,
        )
        handle = KVHandle.create(num_tokens=4, request_id="req-no-retry")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[99])
        assert engine.get_retry_count(tid) == 0

    def test_raises_when_all_retries_exhausted(self) -> None:
        """fail_count > max_retries 时，应抛出 KVTransferFailedError。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _TransientFailBackend(rank=0, mailbox=mailbox, fail_count=5),
            max_retries=2,
            retry_backoff_ms=1.0,
        )
        handle = KVHandle.create(num_tokens=4, request_id="req-retry-fail")

        with pytest.raises(KVTransferFailedError):
            engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[1])

    def test_zero_max_retries_fails_immediately(self) -> None:
        """max_retries=0 时，任何 send 失败直接传播为 KVTransferFailedError。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _TransientFailBackend(rank=0, mailbox=mailbox, fail_count=1),
            max_retries=0,
        )
        handle = KVHandle.create(num_tokens=4, request_id="req-no-retry-fail")

        with pytest.raises(KVTransferFailedError, match="congestion"):
            engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[1])

    def test_retry_count_cleared_on_clear_completed(self) -> None:
        """clear_completed() 后，retry_count 条目应被清理，不应抛出 KeyError。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            max_retries=0,
        )
        handle = KVHandle.create(num_tokens=4, request_id="req-clear")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[7])
        engine.clear_completed()

        with pytest.raises(KeyError):
            engine.get_retry_count(tid)


# ---------------------------------------------------------------------------
# 4. 端到端稳定性/回归基准
# ---------------------------------------------------------------------------


class TestE2ERegressionBenchmark:
    """端对端稳定性回归：多连续传输延迟收敛，指标字段完整，无内存泄漏。"""

    def test_bulk_transfers_complete(self) -> None:
        """连续 20 次跨节点传输全部完成，completed_count 正确。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(_AdvancedMailboxBackend(rank=0, mailbox=mailbox))

        n = 20
        for i in range(n):
            handle = KVHandle.create(num_tokens=8, request_id=f"req-bulk-{i}")
            engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[i])

        assert engine.get_completed_count() == n

    def test_bulk_transfer_latency_bounded(self) -> None:
        """连续 50 次传输，每次延迟应 < 100 ms（in-memory）。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(_AdvancedMailboxBackend(rank=0, mailbox=mailbox))

        for i in range(50):
            handle = KVHandle.create(num_tokens=8, request_id=f"req-latency-{i}")
            tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=bytes(64))
            metrics = engine.get_transfer_metrics(tid)
            assert metrics["latency_ms"] < 100.0, (
                f"Transfer {i} latency too high: {metrics['latency_ms']:.2f} ms"
            )

    def test_metrics_fields_complete(self) -> None:
        """每次传输的 get_transfer_metrics() 应包含所有预期字段。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
            enable_checksum=True,
            max_retries=1,
            retry_backoff_ms=1.0,
        )
        handle = KVHandle.create(num_tokens=8, request_id="req-fields")
        tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=b"field test " * 50)
        metrics = engine.get_transfer_metrics(tid)

        required_keys = {
            "bytes",
            "latency_ms",
            "bandwidth_gbps",
            "shards",
            "compression_ratio",
            "retries",
            "checksum",
            "status",
        }
        assert required_keys.issubset(metrics.keys()), (
            f"Missing keys: {required_keys - set(metrics.keys())}"
        )

    def test_status_completed_for_all_transfers(self) -> None:
        """所有成功传输的 status 应为 'completed'。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(_AdvancedMailboxBackend(rank=0, mailbox=mailbox))

        tids = []
        for i in range(10):
            handle = KVHandle.create(num_tokens=4, request_id=f"req-status-{i}")
            tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[i])
            tids.append(tid)

        for tid in tids:
            assert engine.get_transfer_status(tid) == TransferStatus.COMPLETED

    def test_pending_count_zero_after_all_transfers(self) -> None:
        """所有传输完成后，pending_count 应为 0。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(_AdvancedMailboxBackend(rank=0, mailbox=mailbox))

        for i in range(5):
            handle = KVHandle.create(num_tokens=4, request_id=f"req-pending-{i}")
            engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=[i])

        assert engine.get_pending_count() == 0

    def test_cross_node_bidirectional_transfer(self) -> None:
        """双向传输：rank0→rank1 + rank1→rank2 各自独立完成。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine0 = KVTransferEngine(_AdvancedMailboxBackend(rank=0, mailbox=mailbox))
        engine1 = KVTransferEngine(_AdvancedMailboxBackend(rank=1, mailbox=mailbox))

        # rank0 → rank1
        h01 = KVHandle.create(num_tokens=8, request_id="req-bidir-01")
        recv_buf_01: list[Any] = []
        engine0.transfer_kv_handle(h01, src_rank=0, dst_rank=1, tensor=[10, 20])
        engine1.transfer_kv_handle(h01, src_rank=0, dst_rank=1, tensor=recv_buf_01)
        assert recv_buf_01 == [10, 20]

        # rank1 → rank0
        h10 = KVHandle.create(num_tokens=8, request_id="req-bidir-10")
        recv_buf_10: list[Any] = []
        engine1.transfer_kv_handle(h10, src_rank=1, dst_rank=0, tensor=[30, 40])
        engine0.transfer_kv_handle(h10, src_rank=1, dst_rank=0, tensor=recv_buf_10)
        assert recv_buf_10 == [30, 40]

    def test_compression_stability_over_multiple_payloads(self) -> None:
        """多次启用压缩的传输不累计 is_compressed 状态错误，每次元数据独立。"""
        mailbox: dict[tuple[int, int], list[Any]] = {}
        engine = KVTransferEngine(
            _AdvancedMailboxBackend(rank=0, mailbox=mailbox),
            enable_compression=True,
        )
        for i in range(10):
            payload = bytes(range(i % 256)) * 10
            handle = KVHandle.create(num_tokens=8, request_id=f"req-stable-{i}")
            tid = engine.transfer_kv_handle(handle, src_rank=0, dst_rank=1, tensor=payload)
            meta = engine.get_transfer_metadata(tid)
            assert meta.is_compressed is True
            assert meta.total_bytes > 0
            assert meta.original_size == len(payload)
