"""Agent handshake FSM — manages the 5-phase introduction protocol.

Phases:
  1. TRANSPORT    — Noise_XX encrypted channel established (handled by transport layer)
  2. INTRODUCTION — Agents exchange HELLO messages
  3. VERIFICATION — Mutual nonce challenge-response to prove Ed25519 key ownership
  4. OWNER APPROVAL — Both agents ask their owners for permission
  5. ACTIVE        — Chat approved, normal messaging begins
"""

from __future__ import annotations

import logging
import os
import time
from enum import Enum
from typing import Awaitable, Callable

from a2a_p2p.errors import HandshakeFailed, HandshakeTimeout, PeerRejected
from a2a_p2p.identity import PeerIdentity
from a2a_p2p.peer_store import PeerRecord, PeerStore
from a2a_p2p.protocol.envelope import A2AEnvelope
from a2a_p2p.protocol.types import (
    ApprovalDenyPayload,
    ApprovalGrantPayload,
    ChallengePayload,
    ChallengeResponsePayload,
    HelloPayload,
    MessageType,
)

logger = logging.getLogger(__name__)


class HandshakeState(str, Enum):
    IDLE = "idle"
    TRANSPORT_ESTABLISHED = "transport"
    HELLO_SENT = "hello_sent"
    HELLO_EXCHANGED = "hello_exchanged"
    CHALLENGE_SENT = "challenge_sent"
    VERIFIED = "verified"
    AWAITING_LOCAL_APPROVAL = "awaiting_local"
    AWAITING_REMOTE_APPROVAL = "awaiting_remote"
    ACTIVE = "active"
    REJECTED = "rejected"
    FAILED = "failed"


# Callback: async (peer_id, hello_payload_dict) -> None
ApprovalNeededCallback = Callable[[str, dict], Awaitable[None]]


class AgentHandshake:
    """State machine for a single peer handshake.

    The handshake is driven by incoming messages. Call ``handle_message()``
    for each received envelope. The returned envelope (if any) should be
    sent back to the peer.

    When owner approval is needed, the ``on_approval_needed`` callback fires.
    Call ``owner_decision()`` when the owner responds.
    """

    def __init__(
        self,
        identity: PeerIdentity,
        peer_store: PeerStore,
        on_approval_needed: ApprovalNeededCallback,
        *,
        our_hello: HelloPayload,
        is_initiator: bool = False,
    ) -> None:
        self._identity = identity
        self._peer_store = peer_store
        self._on_approval_needed = on_approval_needed
        self._our_hello = our_hello
        self._is_initiator = is_initiator

        self._state = HandshakeState.TRANSPORT_ESTABLISHED
        self._remote_peer_id: str = ""
        self._peer_hello: HelloPayload | None = None
        self._our_nonce: bytes | None = None
        self._remote_verified = False
        self._remote_grant: ApprovalGrantPayload | None = None
        self._started_at = time.time()

    @property
    def state(self) -> HandshakeState:
        return self._state

    @property
    def remote_peer_id(self) -> str:
        return self._remote_peer_id

    @property
    def peer_hello(self) -> HelloPayload | None:
        return self._peer_hello

    @property
    def is_complete(self) -> bool:
        return self._state == HandshakeState.ACTIVE

    @property
    def is_failed(self) -> bool:
        return self._state in (HandshakeState.REJECTED, HandshakeState.FAILED)

    def set_remote_peer_id(self, peer_id: str) -> None:
        """Set the remote peer ID (from transport-level handshake)."""
        self._remote_peer_id = peer_id

    def start_messages(self) -> list[A2AEnvelope]:
        """Generate the initial messages to send (HELLO, and challenge if initiator)."""
        messages = []

        # Always send HELLO first
        hello_env = A2AEnvelope(
            sender_id=self._identity.peer_id,
            recipient_id=self._remote_peer_id,
            type=MessageType.HELLO,
            payload=self._our_hello.model_dump(),
        )
        messages.append(hello_env)
        self._state = HandshakeState.HELLO_SENT

        return messages

    async def handle_message(self, envelope: A2AEnvelope) -> list[A2AEnvelope]:
        """Process an incoming handshake message.

        Returns a list of response envelopes to send (may be empty).
        """
        msg_type = envelope.type
        responses: list[A2AEnvelope] = []

        if msg_type == MessageType.HELLO:
            responses.extend(self._handle_hello(envelope))

        elif msg_type == MessageType.CHALLENGE:
            responses.extend(self._handle_challenge(envelope))

        elif msg_type == MessageType.CHALLENGE_RESPONSE:
            responses.extend(await self._handle_challenge_response(envelope))

        elif msg_type == MessageType.APPROVAL_GRANT:
            responses.extend(self._handle_approval_grant(envelope))

        elif msg_type == MessageType.APPROVAL_DENY:
            self._handle_approval_deny(envelope)

        return responses

    def _handle_hello(self, envelope: A2AEnvelope) -> list[A2AEnvelope]:
        """Process incoming HELLO."""
        self._peer_hello = HelloPayload.model_validate(envelope.payload)
        if not self._remote_peer_id:
            self._remote_peer_id = envelope.sender_id

        logger.info(
            "Handshake HELLO from %s (%s, owner: %s)",
            self._remote_peer_id,
            self._peer_hello.agent_name,
            self._peer_hello.owner_display_name,
        )

        responses = []

        if self._state == HandshakeState.HELLO_SENT:
            # Both HELLOs exchanged
            self._state = HandshakeState.HELLO_EXCHANGED
        elif self._state == HandshakeState.TRANSPORT_ESTABLISHED:
            # We received their HELLO first — send ours back
            hello_env = A2AEnvelope(
                sender_id=self._identity.peer_id,
                recipient_id=self._remote_peer_id,
                type=MessageType.HELLO,
                payload=self._our_hello.model_dump(),
            )
            responses.append(hello_env)
            self._state = HandshakeState.HELLO_EXCHANGED

        # Initiator sends challenge after HELLO exchange
        if self._state == HandshakeState.HELLO_EXCHANGED and self._is_initiator:
            responses.extend(self._send_challenge())

        return responses

    def _send_challenge(self) -> list[A2AEnvelope]:
        """Generate and send a nonce challenge."""
        self._our_nonce = os.urandom(32)
        challenge = ChallengePayload(
            nonce_hex=self._our_nonce.hex(),
            timestamp=time.time(),
        )
        env = A2AEnvelope(
            sender_id=self._identity.peer_id,
            recipient_id=self._remote_peer_id,
            type=MessageType.CHALLENGE,
            payload=challenge.model_dump(),
        )
        self._state = HandshakeState.CHALLENGE_SENT
        return [env]

    def _handle_challenge(self, envelope: A2AEnvelope) -> list[A2AEnvelope]:
        """Respond to a nonce challenge, then send our own."""
        challenge = ChallengePayload.model_validate(envelope.payload)
        nonce = bytes.fromhex(challenge.nonce_hex)

        # Sign: nonce + our_peer_id
        sign_data = nonce + self._identity.peer_id.encode("utf-8")
        signature = self._identity.sign(sign_data)

        resp = ChallengeResponsePayload(
            nonce_hex=challenge.nonce_hex,
            signature_hex=signature.hex(),
        )
        resp_env = A2AEnvelope(
            sender_id=self._identity.peer_id,
            recipient_id=self._remote_peer_id,
            type=MessageType.CHALLENGE_RESPONSE,
            payload=resp.model_dump(),
        )

        responses = [resp_env]

        # If we haven't sent our own challenge yet, send one
        if self._our_nonce is None:
            responses.extend(self._send_challenge())

        return responses

    async def _handle_challenge_response(
        self, envelope: A2AEnvelope,
    ) -> list[A2AEnvelope]:
        """Verify the peer's challenge response."""
        resp = ChallengeResponsePayload.model_validate(envelope.payload)

        if self._our_nonce is None:
            raise HandshakeFailed(
                self._remote_peer_id, "challenge_response",
                "Received response without sending challenge",
            )

        expected_nonce_hex = self._our_nonce.hex()
        if resp.nonce_hex != expected_nonce_hex:
            raise HandshakeFailed(
                self._remote_peer_id, "challenge_response", "Nonce mismatch",
            )

        # Verify signature
        nonce = self._our_nonce
        sign_data = nonce + self._remote_peer_id.encode("utf-8")
        signature = bytes.fromhex(resp.signature_hex)

        # Get peer's public key from their HELLO or peer store
        peer = self._peer_store.get(self._remote_peer_id)
        if peer and peer.public_key_pem:
            pub_key = PeerIdentity.load_public_key_pem(peer.public_key_pem)
        else:
            # For new peers, we don't have their key yet —
            # trust transport-level verification for now
            logger.info(
                "No stored public key for %s, trusting transport-level verification",
                self._remote_peer_id,
            )
            self._remote_verified = True
            return await self._check_both_verified()

        verified = PeerIdentity.verify(sign_data, signature, pub_key)
        if not verified:
            raise HandshakeFailed(
                self._remote_peer_id, "challenge_response",
                "Signature verification failed",
            )

        logger.info("Challenge-response verified for %s", self._remote_peer_id)
        self._remote_verified = True
        return await self._check_both_verified()

    async def _check_both_verified(self) -> list[A2AEnvelope]:
        """Check if verification is complete and trigger owner approval."""
        if not self._remote_verified:
            return []

        self._state = HandshakeState.VERIFIED

        # Notify owner for approval
        self._state = HandshakeState.AWAITING_LOCAL_APPROVAL
        try:
            hello_dict = self._peer_hello.model_dump() if self._peer_hello else {}
            await self._on_approval_needed(self._remote_peer_id, hello_dict)
        except Exception:
            logger.warning(
                "Failed to notify owner about handshake with %s",
                self._remote_peer_id, exc_info=True,
            )

        return []

    async def owner_decision(
        self,
        approved: bool,
        trust_level: int = 1,
        scope_instructions: str = "",
    ) -> list[A2AEnvelope]:
        """Called when the owner approves or rejects the connection.

        ``scope_instructions`` are LOCAL policy — they are stored in our own
        peer record and never transmitted to the remote peer.
        """
        if self._state not in (
            HandshakeState.AWAITING_LOCAL_APPROVAL,
            HandshakeState.AWAITING_REMOTE_APPROVAL,
        ):
            logger.warning(
                "Owner decision received in unexpected state: %s", self._state,
            )
            return []

        responses = []

        if approved:
            # Only send trust_level — scope_instructions are local policy
            grant = ApprovalGrantPayload(trust_level=trust_level)
            env = A2AEnvelope(
                sender_id=self._identity.peer_id,
                recipient_id=self._remote_peer_id,
                type=MessageType.APPROVAL_GRANT,
                payload=grant.model_dump(),
            )
            responses.append(env)

            # Store local scope for _finalize_approval
            self._local_scope = scope_instructions
            self._local_trust = trust_level

            if self._remote_grant is not None:
                # Remote already approved — both sides done
                self._finalize_approval()
            elif self._state == HandshakeState.AWAITING_LOCAL_APPROVAL:
                self._state = HandshakeState.AWAITING_REMOTE_APPROVAL
            else:
                # Both approved
                self._finalize_approval()
        else:
            deny = ApprovalDenyPayload(reason="Owner declined the connection")
            env = A2AEnvelope(
                sender_id=self._identity.peer_id,
                recipient_id=self._remote_peer_id,
                type=MessageType.APPROVAL_DENY,
                payload=deny.model_dump(),
            )
            responses.append(env)
            self._state = HandshakeState.REJECTED

        return responses

    def _handle_approval_grant(self, envelope: A2AEnvelope) -> list[A2AEnvelope]:
        """Remote owner approved the connection."""
        grant = ApprovalGrantPayload.model_validate(envelope.payload)

        logger.info(
            "Remote owner approved connection with %s (trust=%d)",
            self._remote_peer_id, grant.trust_level,
        )

        if self._state == HandshakeState.AWAITING_REMOTE_APPROVAL:
            # We already approved — both sides done
            self._finalize_approval()
        elif self._state == HandshakeState.AWAITING_LOCAL_APPROVAL:
            # They approved first, we're still waiting for our owner
            # Store their grant for when we approve
            self._remote_grant = grant
        elif self._state == HandshakeState.VERIFIED:
            # They approved before we even asked our owner
            self._remote_grant = grant

        return []

    def _handle_approval_deny(self, envelope: A2AEnvelope) -> None:
        """Remote owner rejected the connection."""
        deny = ApprovalDenyPayload.model_validate(envelope.payload)
        logger.info(
            "Remote owner denied connection with %s: %s",
            self._remote_peer_id, deny.reason,
        )
        self._state = HandshakeState.REJECTED

    def _finalize_approval(self) -> None:
        """Save peer record and transition to ACTIVE.

        Uses ``_local_trust`` and ``_local_scope`` set by ``owner_decision()``
        — these are OUR owner's policy for this peer, never the remote peer's.
        """
        self._state = HandshakeState.ACTIVE

        local_trust = getattr(self, "_local_trust", 1)
        local_scope = getattr(self, "_local_scope", "")

        # Save/update peer record
        existing = self._peer_store.get(self._remote_peer_id)
        hello = self._peer_hello

        record = PeerRecord(
            peer_id=self._remote_peer_id,
            alias=hello.agent_name if hello else existing.alias if existing else "",
            public_key_pem=existing.public_key_pem if existing else "",
            addresses=existing.addresses if existing else [],
            trust_level=max(1, local_trust),  # at least VERIFIED after handshake
            scope_instructions=local_scope,
            description="",  # agent will fill this in later
            capabilities=hello.capabilities if hello else [],
            last_seen=time.time(),
            added_at=existing.added_at if existing else time.time(),
            metadata={
                "owner_display_name": hello.owner_display_name if hello else "",
                "introduction": hello.introduction if hello else "",
                "protocol_version": hello.protocol_version if hello else 1,
            },
        )
        self._peer_store.save(record)

        logger.info(
            "Handshake complete with %s (%s) — trust_level=%d",
            self._remote_peer_id,
            record.alias,
            record.trust_level,
        )
