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

1. TRANSPORT    — Noise_XX encrypted channel + identity binding proof
2. INTRODUCTION — Agents exchange signed HELLO messages (includes Ed25519 pubkey)
3. VERIFICATION — Mutual nonce challenge-response (Ed25519, ALWAYS verified)
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 cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

from p2p.errors import HandshakeFailed
from p2p.identity import PeerIdentity, _derive_peer_id_from_raw
from p2p.peer_store import PeerRecord, PeerStore
from p2p.protocol.envelope import A2AEnvelope
from 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"


ApprovalCallback = Callable[[str, dict], Awaitable[None]]


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

    All envelopes are signed. Challenge verification is mandatory —
    never skipped, even for first-contact peers (uses Ed25519 key from HELLO).
    """

    def __init__(
        self,
        identity: PeerIdentity,
        peer_store: PeerStore,
        on_approval_needed: ApprovalCallback,
        *,
        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
        # Always include our Ed25519 pubkey in HELLO
        self._our_hello.ed25519_pubkey_hex = identity.public_key_raw().hex()
        self._is_initiator = is_initiator
        self._state = HandshakeState.TRANSPORT_ESTABLISHED
        self._remote_peer_id = ""
        self._peer_hello: HelloPayload | None = None
        self._our_nonce: bytes | None = None
        self._remote_grant: ApprovalGrantPayload | None = None
        self._remote_ed25519_pub_hex: str = ""

    @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, pid: str) -> None:
        self._remote_peer_id = pid

    # ── Signed envelope helper ──────────────────────────────────────

    def _signed(self, to: str, t: MessageType, payload: dict) -> A2AEnvelope:
        env = A2AEnvelope(
            sender_id=self._identity.peer_id,
            recipient_id=to,
            type=t,
            payload=payload,
        )
        env.signature_hex = self._identity.sign(env.signing_material()).hex()
        return env

    # ── Public API ──────────────────────────────────────────────────

    def start_messages(self) -> list[A2AEnvelope]:
        env = self._signed(
            self._remote_peer_id, MessageType.HELLO,
            self._our_hello.model_dump(),
        )
        self._state = HandshakeState.HELLO_SENT
        return [env]

    async def handle_message(self, env: A2AEnvelope) -> list[A2AEnvelope]:
        t = env.type
        if t == MessageType.HELLO:
            return self._on_hello(env)
        if t == MessageType.CHALLENGE:
            return self._on_challenge(env)
        if t == MessageType.CHALLENGE_RESPONSE:
            return await self._on_challenge_resp(env)
        if t == MessageType.APPROVAL_GRANT:
            return self._on_grant(env)
        if t == MessageType.APPROVAL_DENY:
            self._state = HandshakeState.REJECTED
        return []

    async def owner_decision(
        self, approved: bool, trust_level: int = 1,
        scope_instructions: str = "",
    ) -> list[A2AEnvelope]:
        if self._state not in (
            HandshakeState.AWAITING_LOCAL_APPROVAL,
            HandshakeState.AWAITING_REMOTE_APPROVAL,
        ):
            return []

        if approved:
            env = self._signed(
                self._remote_peer_id, MessageType.APPROVAL_GRANT,
                ApprovalGrantPayload(
                    trust_level=trust_level,
                    scope_instructions=scope_instructions,
                ).model_dump(),
            )
            if self._remote_grant is not None:
                self._finalize(self._remote_grant.trust_level,
                               self._remote_grant.scope_instructions)
            elif self._state == HandshakeState.AWAITING_LOCAL_APPROVAL:
                self._state = HandshakeState.AWAITING_REMOTE_APPROVAL
            else:
                self._finalize(trust_level, scope_instructions)
            return [env]
        else:
            env = self._signed(
                self._remote_peer_id, MessageType.APPROVAL_DENY,
                ApprovalDenyPayload(reason="Owner declined").model_dump(),
            )
            self._state = HandshakeState.REJECTED
            return [env]

    # ── Internal ────────────────────────────────────────────────────

    def _on_hello(self, env: A2AEnvelope) -> list[A2AEnvelope]:
        self._peer_hello = HelloPayload.model_validate(env.payload)
        if not self._remote_peer_id:
            self._remote_peer_id = env.sender_id

        # Capture Ed25519 pubkey for challenge verification
        self._remote_ed25519_pub_hex = self._peer_hello.ed25519_pubkey_hex

        out: list[A2AEnvelope] = []
        if self._state == HandshakeState.HELLO_SENT:
            self._state = HandshakeState.HELLO_EXCHANGED
        elif self._state == HandshakeState.TRANSPORT_ESTABLISHED:
            out.append(self._signed(
                self._remote_peer_id, MessageType.HELLO,
                self._our_hello.model_dump(),
            ))
            self._state = HandshakeState.HELLO_EXCHANGED
        if self._state == HandshakeState.HELLO_EXCHANGED and self._is_initiator:
            out.extend(self._send_challenge())
        return out

    def _send_challenge(self) -> list[A2AEnvelope]:
        self._our_nonce = os.urandom(32)
        env = self._signed(
            self._remote_peer_id, MessageType.CHALLENGE,
            ChallengePayload(
                nonce_hex=self._our_nonce.hex(), timestamp=time.time(),
            ).model_dump(),
        )
        self._state = HandshakeState.CHALLENGE_SENT
        return [env]

    def _on_challenge(self, env: A2AEnvelope) -> list[A2AEnvelope]:
        ch = ChallengePayload.model_validate(env.payload)
        nonce = bytes.fromhex(ch.nonce_hex)
        sig = self._identity.sign(nonce + self._identity.peer_id.encode())
        resp = self._signed(
            self._remote_peer_id, MessageType.CHALLENGE_RESPONSE,
            ChallengeResponsePayload(
                nonce_hex=ch.nonce_hex, signature_hex=sig.hex(),
            ).model_dump(),
        )
        out = [resp]
        if self._our_nonce is None:
            out.extend(self._send_challenge())
        return out

    async def _on_challenge_resp(self, env: A2AEnvelope) -> list[A2AEnvelope]:
        resp = ChallengeResponsePayload.model_validate(env.payload)

        if not self._our_nonce or resp.nonce_hex != self._our_nonce.hex():
            raise HandshakeFailed(self._remote_peer_id, "challenge", "nonce mismatch")

        # Resolve the peer's Ed25519 public key — NEVER skip verification
        pub: Ed25519PublicKey | None = None

        # Option 1: known peer in store
        peer = self._peer_store.get(self._remote_peer_id)
        if peer and peer.public_key_pem:
            pub = PeerIdentity.load_public_key_pem(peer.public_key_pem)

        # Option 2: first contact — use Ed25519 key from HELLO
        elif self._remote_ed25519_pub_hex:
            try:
                pub_bytes = bytes.fromhex(self._remote_ed25519_pub_hex)
                pub = Ed25519PublicKey.from_public_bytes(pub_bytes)
                # Verify PeerID matches this key
                expected = _derive_peer_id_from_raw(pub_bytes)
                if expected != self._remote_peer_id:
                    raise HandshakeFailed(
                        self._remote_peer_id, "challenge",
                        f"ed25519 key does not match peer_id (expected {expected})",
                    )
            except HandshakeFailed:
                raise
            except Exception as e:
                raise HandshakeFailed(
                    self._remote_peer_id, "challenge",
                    f"invalid ed25519 key in HELLO: {e}",
                ) from e

        if pub is None:
            raise HandshakeFailed(
                self._remote_peer_id, "challenge",
                "no public key available — peer must include ed25519_pubkey_hex in HELLO",
            )

        # Verify — ALWAYS
        sign_data = self._our_nonce + self._remote_peer_id.encode()
        if not PeerIdentity.verify(sign_data, bytes.fromhex(resp.signature_hex), pub):
            raise HandshakeFailed(self._remote_peer_id, "challenge", "bad signature")

        logger.info("Challenge verified for %s", self._remote_peer_id)

        # Move to owner approval
        self._state = HandshakeState.AWAITING_LOCAL_APPROVAL
        try:
            await self._on_approval_needed(
                self._remote_peer_id,
                self._peer_hello.model_dump() if self._peer_hello else {},
            )
        except Exception:
            logger.warning("Approval callback failed", exc_info=True)
        return []

    def _on_grant(self, env: A2AEnvelope) -> list[A2AEnvelope]:
        grant = ApprovalGrantPayload.model_validate(env.payload)
        if self._state == HandshakeState.AWAITING_REMOTE_APPROVAL:
            self._finalize(grant.trust_level, grant.scope_instructions)
        else:
            self._remote_grant = grant
        return []

    def _finalize(self, remote_trust: int, remote_scope: str) -> None:
        self._state = HandshakeState.ACTIVE
        existing = self._peer_store.get(self._remote_peer_id)
        h = self._peer_hello

        # Derive public key PEM for first-contact peers
        public_key_pem = ""
        if existing and existing.public_key_pem:
            public_key_pem = existing.public_key_pem
        elif self._remote_ed25519_pub_hex:
            try:
                pub_bytes = bytes.fromhex(self._remote_ed25519_pub_hex)
                pub = Ed25519PublicKey.from_public_bytes(pub_bytes)
                public_key_pem = pub.public_bytes(
                    Encoding.PEM, PublicFormat.SubjectPublicKeyInfo,
                ).decode("ascii")
            except Exception:
                logger.warning("Failed to derive PEM from HELLO pubkey")

        self._peer_store.save(PeerRecord(
            peer_id=self._remote_peer_id,
            alias=h.agent_name if h else (existing.alias if existing else ""),
            public_key_pem=public_key_pem,
            addresses=existing.addresses if existing else [],
            trust_level=max(1, remote_trust),
            scope_instructions=remote_scope,
            capabilities=h.capabilities if h else [],
            last_seen=time.time(),
            added_at=existing.added_at if existing else time.time(),
            metadata={
                "owner_display_name": h.owner_display_name if h else "",
                "introduction": h.introduction if h else "",
            },
        ))
        logger.info("Handshake complete: %s (pubkey stored)", self._remote_peer_id)
