"""Tests for AnycastAgent."""

import asyncio
import json
import pytest
import websockets
import websockets.server
from anycast_agent import (
    AnycastAgent,
    ConnectionState,
    AnycastError,
    AnycastAuthError,
    AnycastTimeoutError,
)


@pytest.fixture
async def echo_server():
    """WebSocket server that echoes registration."""

    async def handler(ws):
        async for raw in ws:
            msg = json.loads(raw)
            if msg["type"] == "register":
                await ws.send(json.dumps({
                    "type": "registered",
                    "agentId": "agent-123",
                    "tenantId": "tenant-456",
                    "name": "Test Agent",
                }))
            elif msg["type"] == "heartbeat":
                await ws.send(json.dumps({
                    "type": "heartbeat_ack",
                    "timestamp": 1234567890,
                }))

    server = await websockets.serve(handler, "localhost", 0)
    port = server.sockets[0].getsockname()[1]
    yield f"ws://localhost:{port}"
    server.close()
    await server.wait_closed()


@pytest.fixture
async def auth_fail_server():
    """Server that rejects registration."""

    async def handler(ws):
        async for raw in ws:
            msg = json.loads(raw)
            if msg["type"] == "register":
                await ws.send(json.dumps({
                    "type": "error",
                    "code": "AUTH_FAILED",
                    "message": "Invalid token",
                    "fatal": True,
                }))

    server = await websockets.serve(handler, "localhost", 0)
    port = server.sockets[0].getsockname()[1]
    yield f"ws://localhost:{port}"
    server.close()
    await server.wait_closed()


@pytest.fixture
async def silent_server():
    """Server that accepts but never responds."""

    async def handler(ws):
        await asyncio.sleep(60)

    server = await websockets.serve(handler, "localhost", 0)
    port = server.sockets[0].getsockname()[1]
    yield f"ws://localhost:{port}"
    server.close()
    await server.wait_closed()


@pytest.fixture
async def peer_server():
    """Server that sends peer status after registration."""

    async def handler(ws):
        async for raw in ws:
            msg = json.loads(raw)
            if msg["type"] == "register":
                await ws.send(json.dumps({
                    "type": "registered",
                    "agentId": "agent-1",
                    "tenantId": "tenant-1",
                    "name": "Agent 1",
                }))
                await asyncio.sleep(0.05)
                await ws.send(json.dumps({
                    "type": "agent_status",
                    "agentId": "agent-2",
                    "status": "online",
                }))

    server = await websockets.serve(handler, "localhost", 0)
    port = server.sockets[0].getsockname()[1]
    yield f"ws://localhost:{port}"
    server.close()
    await server.wait_closed()


class TestAnycastAgent:
    @pytest.mark.asyncio
    async def test_initial_state(self, echo_server):
        agent = AnycastAgent(rendezvous_url=echo_server, token="test")
        assert agent.connection_state == ConnectionState.DISCONNECTED
        assert agent.connected is False
        assert agent.agent_id is None

    @pytest.mark.asyncio
    async def test_connect_and_register(self, echo_server):
        agent = AnycastAgent(
            rendezvous_url=echo_server,
            token="test",
            reconnect=False,
        )

        events = []
        agent.on("connect", lambda info: events.append(("connect", info.agent_id)))

        info = await agent.connect()

        assert agent.connected is True
        assert agent.agent_id == "agent-123"
        assert agent.tenant_id == "tenant-456"
        assert info.agent_id == "agent-123"
        assert info.name == "Test Agent"
        assert len(events) == 1
        assert events[0] == ("connect", "agent-123")

        await agent.disconnect()
        assert agent.connected is False
        assert agent.agent_id is None

    @pytest.mark.asyncio
    async def test_auth_failure(self, auth_fail_server):
        agent = AnycastAgent(
            rendezvous_url=auth_fail_server,
            token="bad-token",
            reconnect=False,
        )

        with pytest.raises(AnycastAuthError, match="Invalid token"):
            await agent.connect()

        assert agent.connected is False

    @pytest.mark.asyncio
    async def test_timeout(self, silent_server):
        agent = AnycastAgent(
            rendezvous_url=silent_server,
            token="test",
            timeout=0.5,
            reconnect=False,
        )

        with pytest.raises(AnycastTimeoutError):
            await agent.connect()

    @pytest.mark.asyncio
    async def test_peer_status_event(self, peer_server):
        agent = AnycastAgent(
            rendezvous_url=peer_server,
            token="test",
            reconnect=False,
        )

        peers = []
        agent.on("peer", lambda aid, status: peers.append((aid, status)))

        await agent.connect()
        await asyncio.sleep(0.2)

        assert len(peers) == 1
        assert peers[0] == ("agent-2", "online")

        await agent.disconnect()

    @pytest.mark.asyncio
    async def test_heartbeat(self, echo_server):
        agent = AnycastAgent(
            rendezvous_url=echo_server,
            token="test",
            heartbeat_interval=0.1,
            reconnect=False,
        )

        await agent.connect()
        await asyncio.sleep(0.35)
        # If heartbeat failed the connection would have dropped
        assert agent.connected is True
        await agent.disconnect()

    @pytest.mark.asyncio
    async def test_request_connection(self, echo_server):
        agent = AnycastAgent(
            rendezvous_url=echo_server,
            token="test",
            reconnect=False,
        )

        await agent.connect()

        session_id = agent.request_connection("agent-2")
        assert isinstance(session_id, str)
        assert len(session_id) == 32  # hex of 16 bytes

        await agent.disconnect()

    @pytest.mark.asyncio
    async def test_not_connected_error(self, echo_server):
        agent = AnycastAgent(
            rendezvous_url=echo_server,
            token="test",
            reconnect=False,
        )

        with pytest.raises(AnycastError, match="Not connected"):
            agent.request_connection("agent-2")

    @pytest.mark.asyncio
    async def test_decorator_style(self, echo_server):
        agent = AnycastAgent(
            rendezvous_url=echo_server,
            token="test",
            reconnect=False,
        )

        connected = []

        @agent.on_connect
        async def handle_connect(info):
            connected.append(info.agent_id)

        await agent.connect()
        assert connected == ["agent-123"]
        await agent.disconnect()
