"""Tests for the agent handshake FSM."""

import asyncio
import tempfile
from pathlib import Path

import pytest

from a2a_p2p.handshake import AgentHandshake, HandshakeState
from a2a_p2p.identity import PeerIdentity
from a2a_p2p.peer_store import PeerStore
from a2a_p2p.protocol.envelope import A2AEnvelope
from a2a_p2p.protocol.types import HelloPayload, MessageType


def _make_hello(name: str) -> HelloPayload:
    return HelloPayload(
        agent_name=name,
        owner_display_name=f"{name}'s owner",
        capabilities=["test"],
        introduction=f"I am {name}",
    )


class TestHandshakeFSM:
    """Test the handshake state machine with two agents."""

    def _setup(self):
        """Create two identities and peer stores."""
        tmp = tempfile.mkdtemp()
        id_a = PeerIdentity.generate()
        id_b = PeerIdentity.generate()
        store_a = PeerStore(Path(tmp) / "peers_a.json")
        store_b = PeerStore(Path(tmp) / "peers_b.json")
        return id_a, id_b, store_a, store_b

    @pytest.mark.asyncio
    async def test_full_handshake_flow(self):
        """Two agents go through HELLO → challenge → verify → approve."""
        id_a, id_b, store_a, store_b = self._setup()

        approvals_a = []
        approvals_b = []

        async def on_approval_a(pid, hello):
            approvals_a.append((pid, hello))

        async def on_approval_b(pid, hello):
            approvals_b.append((pid, hello))

        hs_a = AgentHandshake(
            identity=id_a,
            peer_store=store_a,
            on_approval_needed=on_approval_a,
            our_hello=_make_hello("Jarvis"),
            is_initiator=True,
        )
        hs_a.set_remote_peer_id(id_b.peer_id)

        hs_b = AgentHandshake(
            identity=id_b,
            peer_store=store_b,
            on_approval_needed=on_approval_b,
            our_hello=_make_hello("Nova"),
            is_initiator=False,
        )
        hs_b.set_remote_peer_id(id_a.peer_id)

        # Start: A sends HELLO
        msgs_a = hs_a.start_messages()
        assert len(msgs_a) == 1
        assert msgs_a[0].type == MessageType.HELLO
        assert hs_a.state == HandshakeState.HELLO_SENT

        # B receives HELLO, sends back HELLO
        msgs_b = await hs_b.handle_message(msgs_a[0])
        # B should respond with HELLO (since it's in TRANSPORT_ESTABLISHED state)
        hello_from_b = [m for m in msgs_b if m.type == MessageType.HELLO]
        assert len(hello_from_b) == 1

        # A receives HELLO from B → triggers challenge (since A is initiator)
        msgs_a2 = await hs_a.handle_message(hello_from_b[0])
        challenges_from_a = [m for m in msgs_a2 if m.type == MessageType.CHALLENGE]
        assert len(challenges_from_a) == 1
        assert hs_a.state == HandshakeState.CHALLENGE_SENT

        # B receives challenge → sends response + own challenge
        msgs_b2 = await hs_b.handle_message(challenges_from_a[0])
        responses_from_b = [m for m in msgs_b2 if m.type == MessageType.CHALLENGE_RESPONSE]
        challenges_from_b = [m for m in msgs_b2 if m.type == MessageType.CHALLENGE]
        assert len(responses_from_b) == 1
        assert len(challenges_from_b) == 1

        # A receives B's challenge response → verifies
        msgs_a3 = await hs_a.handle_message(responses_from_b[0])

        # A receives B's challenge → sends response
        msgs_a4 = await hs_a.handle_message(challenges_from_b[0])
        responses_from_a = [m for m in msgs_a4 if m.type == MessageType.CHALLENGE_RESPONSE]
        assert len(responses_from_a) == 1

        # B receives A's challenge response → verifies
        await hs_b.handle_message(responses_from_a[0])

        # Both should now be awaiting local approval
        assert hs_a.state == HandshakeState.AWAITING_LOCAL_APPROVAL
        assert hs_b.state == HandshakeState.AWAITING_LOCAL_APPROVAL
        assert len(approvals_a) == 1
        assert len(approvals_b) == 1

        # Owner A approves
        approve_msgs_a = await hs_a.owner_decision(True, trust_level=2, scope_instructions="be helpful")
        assert len(approve_msgs_a) == 1
        assert approve_msgs_a[0].type == MessageType.APPROVAL_GRANT
        # Scope instructions must NOT appear in the wire payload
        assert "scope_instructions" not in approve_msgs_a[0].payload
        assert hs_a.state == HandshakeState.AWAITING_REMOTE_APPROVAL

        # B receives A's approval
        await hs_b.handle_message(approve_msgs_a[0])

        # Owner B approves
        approve_msgs_b = await hs_b.owner_decision(True, trust_level=2, scope_instructions="research only")
        assert len(approve_msgs_b) == 1
        assert approve_msgs_b[0].type == MessageType.APPROVAL_GRANT
        assert hs_b.state == HandshakeState.ACTIVE

        # A receives B's approval
        await hs_a.handle_message(approve_msgs_b[0])
        assert hs_a.state == HandshakeState.ACTIVE

        # Verify peer records saved
        rec_a = store_a.get(id_b.peer_id)
        assert rec_a is not None
        assert rec_a.alias == "Nova"
        assert rec_a.trust_level >= 1

        rec_b = store_b.get(id_a.peer_id)
        assert rec_b is not None
        assert rec_b.alias == "Jarvis"

        # Scope instructions must stay local — never leaked to the remote peer
        assert rec_a.scope_instructions == "be helpful"      # A's own policy for B
        assert rec_b.scope_instructions == "research only"   # B's own policy for A
        # NOT the other way around (which was the bug)

    @pytest.mark.asyncio
    async def test_rejection(self):
        """One side rejects the connection."""
        id_a, id_b, store_a, store_b = self._setup()

        async def noop(pid, hello):
            pass

        hs_a = AgentHandshake(
            identity=id_a, peer_store=store_a, on_approval_needed=noop,
            our_hello=_make_hello("Jarvis"), is_initiator=True,
        )
        hs_a.set_remote_peer_id(id_b.peer_id)
        # Manually move to awaiting local approval
        hs_a._state = HandshakeState.AWAITING_LOCAL_APPROVAL

        reject_msgs = await hs_a.owner_decision(False)
        assert len(reject_msgs) == 1
        assert reject_msgs[0].type == MessageType.APPROVAL_DENY
        assert hs_a.state == HandshakeState.REJECTED

    @pytest.mark.asyncio
    async def test_hello_exchange(self):
        """Verify HELLO payloads are captured correctly."""
        id_a, id_b, store_a, store_b = self._setup()

        async def noop(pid, hello):
            pass

        hs_b = AgentHandshake(
            identity=id_b, peer_store=store_b, on_approval_needed=noop,
            our_hello=_make_hello("Nova"), is_initiator=False,
        )
        hs_b.set_remote_peer_id(id_a.peer_id)

        hello_env = A2AEnvelope(
            sender_id=id_a.peer_id,
            recipient_id=id_b.peer_id,
            type=MessageType.HELLO,
            payload=_make_hello("Jarvis").model_dump(),
        )

        await hs_b.handle_message(hello_env)

        assert hs_b.peer_hello is not None
        assert hs_b.peer_hello.agent_name == "Jarvis"
        assert hs_b.peer_hello.owner_display_name == "Jarvis's owner"
