"""Tests for the agent handshake FSM — security hardened."""

import tempfile
from pathlib import Path
import pytest
from p2p.identity import PeerIdentity
from p2p.peer_store import PeerStore
from p2p.handshake import AgentHandshake, HandshakeState
from p2p.errors import HandshakeFailed
from p2p.protocol.types import HelloPayload, MessageType


def _hello(name):
    return HelloPayload(agent_name=name, owner_display_name=f"{name} owner", capabilities=["test"])


class TestHandshake:
    @pytest.mark.asyncio
    async def test_full_flow(self):
        """Full HELLO → challenge → verify → approve with Ed25519 pubkey in HELLO."""
        tmp = tempfile.mkdtemp()
        ia, ib = PeerIdentity.generate(), PeerIdentity.generate()
        sa, sb = PeerStore(Path(tmp)/"a.json"), PeerStore(Path(tmp)/"b.json")
        approvals = []

        async def cb(pid, h):
            approvals.append(pid)

        a = AgentHandshake(ia, sa, cb, our_hello=_hello("A"), is_initiator=True)
        a.set_remote_peer_id(ib.peer_id)
        b = AgentHandshake(ib, sb, cb, our_hello=_hello("B"), is_initiator=False)
        b.set_remote_peer_id(ia.peer_id)

        # HELLO exchange
        m = a.start_messages()
        assert m[0].signature_hex  # All messages signed
        r = await b.handle_message(m[0])
        hellos = [x for x in r if x.type == MessageType.HELLO]
        assert hellos[0].signature_hex
        r2 = await a.handle_message(hellos[0])

        # Challenge exchange
        chals = [x for x in r2 if x.type == MessageType.CHALLENGE]
        r3 = await b.handle_message(chals[0])
        resps = [x for x in r3 if x.type == MessageType.CHALLENGE_RESPONSE]
        b_chals = [x for x in r3 if x.type == MessageType.CHALLENGE]
        await a.handle_message(resps[0])
        r4 = await a.handle_message(b_chals[0])
        a_resps = [x for x in r4 if x.type == MessageType.CHALLENGE_RESPONSE]
        await b.handle_message(a_resps[0])

        assert a.state == HandshakeState.AWAITING_LOCAL_APPROVAL
        assert b.state == HandshakeState.AWAITING_LOCAL_APPROVAL
        assert len(approvals) == 2

        # Approvals
        ga = await a.owner_decision(True, 2, "scope a")
        await b.handle_message(ga[0])
        gb = await b.owner_decision(True, 2, "scope b")
        await a.handle_message(gb[0])

        assert a.state == HandshakeState.ACTIVE
        assert b.state == HandshakeState.ACTIVE

        # Verify peer records saved with public key
        rec_a = sa.get(ib.peer_id)
        assert rec_a.alias == "B"
        assert rec_a.public_key_pem  # Key stored on first contact

        rec_b = sb.get(ia.peer_id)
        assert rec_b.alias == "A"
        assert rec_b.public_key_pem

    @pytest.mark.asyncio
    async def test_first_contact_verifies_challenge(self):
        """Even with no peer store entry, challenge is verified via HELLO pubkey."""
        tmp = tempfile.mkdtemp()
        ia, ib = PeerIdentity.generate(), PeerIdentity.generate()
        sa, sb = PeerStore(Path(tmp)/"a.json"), PeerStore(Path(tmp)/"b.json")

        async def noop(p, h): pass

        a = AgentHandshake(ia, sa, noop, our_hello=_hello("A"), is_initiator=True)
        a.set_remote_peer_id(ib.peer_id)
        b = AgentHandshake(ib, sb, noop, our_hello=_hello("B"), is_initiator=False)
        b.set_remote_peer_id(ia.peer_id)

        # Run through HELLO + challenge
        m = a.start_messages()
        r = await b.handle_message(m[0])
        hellos = [x for x in r if x.type == MessageType.HELLO]
        r2 = await a.handle_message(hellos[0])
        chals = [x for x in r2 if x.type == MessageType.CHALLENGE]
        r3 = await b.handle_message(chals[0])
        resps = [x for x in r3 if x.type == MessageType.CHALLENGE_RESPONSE]

        # A should verify B's challenge response using the HELLO pubkey
        # (B is NOT in A's peer store)
        assert sa.get(ib.peer_id) is None
        await a.handle_message(resps[0])  # Should NOT raise
        assert a.state == HandshakeState.AWAITING_LOCAL_APPROVAL

    @pytest.mark.asyncio
    async def test_wrong_pubkey_in_hello_fails(self):
        """If HELLO contains a different Ed25519 key, challenge fails."""
        tmp = tempfile.mkdtemp()
        ia, ib = PeerIdentity.generate(), PeerIdentity.generate()
        sa, sb = PeerStore(Path(tmp)/"a.json"), PeerStore(Path(tmp)/"b.json")

        async def noop(p, h): pass

        # B will use a HELLO with a DIFFERENT Ed25519 key (not matching its peer_id)
        fake_identity = PeerIdentity.generate()
        bad_hello = _hello("B")
        bad_hello.ed25519_pubkey_hex = fake_identity.public_key_raw().hex()

        a = AgentHandshake(ia, sa, noop, our_hello=_hello("A"), is_initiator=True)
        a.set_remote_peer_id(ib.peer_id)
        b = AgentHandshake(ib, sb, noop, our_hello=_hello("B"), is_initiator=False)
        b.set_remote_peer_id(ia.peer_id)

        # Exchange HELLOs normally
        m = a.start_messages()
        r = await b.handle_message(m[0])
        hellos = [x for x in r if x.type == MessageType.HELLO]

        # Tamper: inject bad pubkey into the HELLO A receives
        from p2p.protocol.envelope import A2AEnvelope
        tampered = A2AEnvelope.model_validate(hellos[0].model_dump())
        payload = dict(tampered.payload)
        payload["ed25519_pubkey_hex"] = fake_identity.public_key_raw().hex()
        tampered.payload = payload
        r2 = await a.handle_message(tampered)

        chals = [x for x in r2 if x.type == MessageType.CHALLENGE]
        r3 = await b.handle_message(chals[0])
        resps = [x for x in r3 if x.type == MessageType.CHALLENGE_RESPONSE]

        # A should reject: the HELLO pubkey doesn't match B's peer_id
        with pytest.raises(HandshakeFailed, match="ed25519 key does not match"):
            await a.handle_message(resps[0])

    @pytest.mark.asyncio
    async def test_reject(self):
        async def noop(p, h): pass
        a = AgentHandshake(PeerIdentity.generate(), PeerStore("/dev/null"), noop,
                           our_hello=_hello("A"), is_initiator=True)
        a._state = HandshakeState.AWAITING_LOCAL_APPROVAL
        m = await a.owner_decision(False)
        assert m[0].type == MessageType.APPROVAL_DENY
        assert m[0].signature_hex  # Signed
        assert a.state == HandshakeState.REJECTED

    @pytest.mark.asyncio
    async def test_no_pubkey_in_hello_fails(self):
        """If peer sends HELLO without ed25519_pubkey_hex, challenge fails."""
        tmp = tempfile.mkdtemp()
        ia, ib = PeerIdentity.generate(), PeerIdentity.generate()
        sa = PeerStore(Path(tmp)/"a.json")

        async def noop(p, h): pass

        a = AgentHandshake(ia, sa, noop, our_hello=_hello("A"), is_initiator=True)
        a.set_remote_peer_id(ib.peer_id)

        # Simulate HELLO without pubkey
        from p2p.protocol.envelope import A2AEnvelope
        from p2p.protocol.types import ChallengeResponsePayload
        import os

        # Set state as if we received a HELLO without pubkey
        a._state = HandshakeState.CHALLENGE_SENT
        a._our_nonce = os.urandom(32)
        a._remote_ed25519_pub_hex = ""  # No pubkey

        # B signs the challenge response correctly
        sig = ib.sign(a._our_nonce + ib.peer_id.encode())
        resp_env = A2AEnvelope(
            sender_id=ib.peer_id, type=MessageType.CHALLENGE_RESPONSE,
            payload=ChallengeResponsePayload(
                nonce_hex=a._our_nonce.hex(), signature_hex=sig.hex(),
            ).model_dump(),
        )

        with pytest.raises(HandshakeFailed, match="no public key available"):
            await a.handle_message(resp_env)
