"""Tests for comm_types module (issue #16 / #17)"""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from sagellm_protocol import (
    CollectiveOp,
    CollectiveRequest,
    CollectiveResponse,
    CommBackendType,
    CommGroupInfo,
    CommOperationMetrics,
    CommStatus,
    CommTopologyDescriptor,
    P2PRecvRequest,
    P2PRecvResponse,
    P2PSendRequest,
    P2PSendResponse,
    ReduceOp,
    TopologyType,
)


class TestCommBackendType:
    def test_all_backends_exist(self) -> None:
        assert CommBackendType.NCCL == "nccl"
        assert CommBackendType.HCCL == "hccl"
        assert CommBackendType.GLOO == "gloo"
        assert CommBackendType.MPI == "mpi"
        assert CommBackendType.CUSTOM == "custom"


class TestReduceOp:
    def test_all_reduce_ops(self) -> None:
        assert ReduceOp.SUM == "sum"
        assert ReduceOp.MAX == "max"
        assert ReduceOp.MIN == "min"
        assert ReduceOp.PROD == "prod"
        assert ReduceOp.AVG == "avg"
        assert ReduceOp.BAND == "band"
        assert ReduceOp.BOR == "bor"
        assert ReduceOp.BXOR == "bxor"


class TestCollectiveOp:
    def test_all_collective_ops(self) -> None:
        assert CollectiveOp.ALL_REDUCE == "all_reduce"
        assert CollectiveOp.ALL_GATHER == "all_gather"
        assert CollectiveOp.REDUCE_SCATTER == "reduce_scatter"
        assert CollectiveOp.BROADCAST == "broadcast"
        assert CollectiveOp.REDUCE == "reduce"
        assert CollectiveOp.ALL_TO_ALL == "all_to_all"
        assert CollectiveOp.BARRIER == "barrier"
        assert CollectiveOp.SEND == "send"
        assert CollectiveOp.RECV == "recv"


class TestCommGroupInfo:
    def test_valid_comm_group(self) -> None:
        group = CommGroupInfo(
            group_id="group-tp-0",
            backend=CommBackendType.NCCL,
            world_size=4,
            ranks=[0, 1, 2, 3],
            root_rank=0,
            timeout_ms=30000.0,
        )
        assert group.group_id == "group-tp-0"
        assert group.backend == CommBackendType.NCCL
        assert group.world_size == 4
        assert group.ranks == [0, 1, 2, 3]
        assert group.root_rank == 0
        assert group.timeout_ms == 30000.0

    def test_world_size_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            CommGroupInfo(
                group_id="g1",
                backend=CommBackendType.NCCL,
                world_size=0,
                ranks=[],
            )

    def test_default_timeout(self) -> None:
        group = CommGroupInfo(
            group_id="g1",
            backend=CommBackendType.GLOO,
            world_size=2,
            ranks=[0, 1],
        )
        assert group.timeout_ms == 30000.0
        assert group.root_rank == 0

    def test_metadata_default_empty(self) -> None:
        group = CommGroupInfo(
            group_id="g1",
            backend=CommBackendType.HCCL,
            world_size=1,
            ranks=[0],
        )
        assert group.metadata == {}


class TestCollectiveRequest:
    def test_valid_all_reduce_request(self) -> None:
        req = CollectiveRequest(
            request_id="req-001",
            trace_id="trace-001",
            group_id="group-tp-0",
            op=CollectiveOp.ALL_REDUCE,
            reduce_op=ReduceOp.SUM,
        )
        assert req.op == CollectiveOp.ALL_REDUCE
        assert req.reduce_op == ReduceOp.SUM
        assert req.async_op is False

    def test_broadcast_with_root(self) -> None:
        req = CollectiveRequest(
            request_id="req-002",
            trace_id="trace-002",
            group_id="group-pp-0",
            op=CollectiveOp.BROADCAST,
            root_rank=0,
        )
        assert req.root_rank == 0

    def test_async_op_flag(self) -> None:
        req = CollectiveRequest(
            request_id="req-003",
            trace_id="trace-003",
            group_id="group-tp-0",
            op=CollectiveOp.ALL_GATHER,
            async_op=True,
        )
        assert req.async_op is True


class TestCollectiveResponse:
    def test_valid_completed_response(self) -> None:
        resp = CollectiveResponse(
            request_id="req-001",
            trace_id="trace-001",
            op=CollectiveOp.ALL_REDUCE,
            status=CommStatus.COMPLETED,
            elapsed_ms=5.2,
            bytes_transferred=1024,
        )
        assert resp.status == CommStatus.COMPLETED
        assert resp.elapsed_ms == 5.2
        assert resp.bytes_transferred == 1024
        assert resp.error_message is None

    def test_failed_response_with_error(self) -> None:
        resp = CollectiveResponse(
            request_id="req-002",
            trace_id="trace-002",
            op=CollectiveOp.ALL_REDUCE,
            status=CommStatus.FAILED,
            error_message="NCCL error: ncclSystemError",
        )
        assert resp.status == CommStatus.FAILED
        assert "NCCL error" in (resp.error_message or "")

    def test_default_values(self) -> None:
        resp = CollectiveResponse(
            request_id="req-003",
            trace_id="trace-003",
            op=CollectiveOp.BARRIER,
            status=CommStatus.COMPLETED,
        )
        assert resp.elapsed_ms == 0.0
        assert resp.bytes_transferred == 0


class TestP2PSendRequest:
    def test_valid_send_request(self) -> None:
        req = P2PSendRequest(
            request_id="p2p-001",
            trace_id="trace-001",
            group_id="group-tp-0",
            src_rank=0,
            dst_rank=1,
            tag=42,
            nbytes=8192,
        )
        assert req.src_rank == 0
        assert req.dst_rank == 1
        assert req.tag == 42
        assert req.nbytes == 8192

    def test_nbytes_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            P2PSendRequest(
                request_id="p2p-001",
                trace_id="trace-001",
                group_id="g1",
                src_rank=0,
                dst_rank=1,
                nbytes=0,
            )


class TestP2PSendResponse:
    def test_send_response(self) -> None:
        resp = P2PSendResponse(
            request_id="p2p-001",
            trace_id="trace-001",
            status=CommStatus.COMPLETED,
            elapsed_ms=1.5,
            bytes_sent=8192,
        )
        assert resp.status == CommStatus.COMPLETED
        assert resp.bytes_sent == 8192


class TestP2PRecvRequest:
    def test_valid_recv_request(self) -> None:
        req = P2PRecvRequest(
            request_id="p2p-002",
            trace_id="trace-002",
            group_id="group-tp-0",
            src_rank=0,
            dst_rank=1,
            expected_nbytes=8192,
        )
        assert req.expected_nbytes == 8192

    def test_expected_nbytes_must_be_positive(self) -> None:
        with pytest.raises(ValidationError):
            P2PRecvRequest(
                request_id="p2p-002",
                trace_id="trace-002",
                group_id="g1",
                src_rank=0,
                dst_rank=1,
                expected_nbytes=0,
            )


class TestP2PRecvResponse:
    def test_recv_response(self) -> None:
        resp = P2PRecvResponse(
            request_id="p2p-002",
            trace_id="trace-002",
            status=CommStatus.COMPLETED,
            elapsed_ms=2.1,
            bytes_received=8192,
        )
        assert resp.status == CommStatus.COMPLETED
        assert resp.bytes_received == 8192


class TestCommOperationMetrics:
    def test_valid_metrics(self) -> None:
        m = CommOperationMetrics(
            op=CollectiveOp.ALL_REDUCE,
            group_id="group-tp-0",
            elapsed_ms=3.5,
            bytes_transferred=16384,
            bandwidth_gbps=4.6,
        )
        assert m.op == CollectiveOp.ALL_REDUCE
        assert m.bandwidth_gbps == 4.6


class TestCommTopologyDescriptor:
    def test_valid_topology(self) -> None:
        topo = CommTopologyDescriptor(
            topology_type=TopologyType.RING,
            num_nodes=4,
            gpus_per_node=8,
            intra_node_bandwidth_gbps=600.0,
            inter_node_bandwidth_gbps=25.0,
            backend=CommBackendType.NCCL,
        )
        assert topo.topology_type == TopologyType.RING
        assert topo.num_nodes == 4
        assert topo.inter_node_bandwidth_gbps == 25.0

    def test_topology_types_exist(self) -> None:
        for t in [
            TopologyType.RING,
            TopologyType.TREE,
            TopologyType.MESH,
            TopologyType.HYBRID,
            TopologyType.FULLY_CONNECTED,
        ]:
            desc = CommTopologyDescriptor(
                topology_type=t,
                num_nodes=2,
                backend=CommBackendType.NCCL,
            )
            assert desc.topology_type == t
