"""A2ANode — main P2P runtime. Manages connections, handshakes, routing."""

from __future__ import annotations

import asyncio
import logging
from collections import deque
from typing import Awaitable, Callable

from p2p.config import A2AConfig
from p2p.errors import PeerUnreachable
from p2p.handshake import AgentHandshake, HandshakeState
from p2p.identity import PeerIdentity
from p2p.peer_store import PeerStore
from p2p.protocol import A2AEnvelope, MessageType, decode, encode
from p2p.protocol.types import HelloPayload, StatusPayload
from p2p.transport.base import Connection
from p2p.transport.websocket import WSTransport

logger = logging.getLogger(__name__)

MessageHandler = Callable[[A2AEnvelope], Awaitable[None]]
ApprovalHandler = Callable[[str, dict], Awaitable[None]]

_HS_TYPES = {
    MessageType.HELLO.value, MessageType.CHALLENGE.value,
    MessageType.CHALLENGE_RESPONSE.value, MessageType.APPROVAL_REQUEST.value,
    MessageType.APPROVAL_GRANT.value, MessageType.APPROVAL_DENY.value,
}


class A2ANode:
    """The P2P node runtime.

    Usage::

        node = A2ANode(identity, peer_store, config, agent_name="Jarvis")
        node.on_message(handler)
        node.on_approval_needed(approval_handler)
        await node.start()
        await node.send(peer_id, MessageType.TEXT, {"text": "hi"})
        await node.stop()
    """

    def __init__(
        self,
        identity: PeerIdentity,
        peer_store: PeerStore,
        config: A2AConfig,
        *,
        agent_name: str = "",
        owner_display_name: str = "",
        capabilities: list[str] | None = None,
        introduction: str = "",
    ) -> None:
        self._identity = identity
        self._peer_store = peer_store
        self._config = config
        self._hello = HelloPayload(
            agent_name=agent_name or "Agent",
            owner_display_name=owner_display_name,
            capabilities=capabilities or [],
            introduction=introduction,
        )
        self._transport = WSTransport(identity)
        self._connections: dict[str, Connection] = {}
        self._handshakes: dict[str, AgentHandshake] = {}
        self._msg_handlers: list[MessageHandler] = []
        self._approval_handler: ApprovalHandler | None = None
        self._tasks: dict[str, asyncio.Task] = {}
        self._listener: asyncio.Task | None = None
        self._reconnector: asyncio.Task | None = None
        self._running = False
        # Message deduplication ring buffer
        self._seen_set: set[str] = set()
        self._seen_queue: deque[str] = deque(maxlen=10000)

    @property
    def peer_id(self) -> str:
        return self._identity.peer_id

    @property
    def connections(self) -> dict[str, Connection]:
        return dict(self._connections)

    @property
    def pending_handshakes(self) -> dict[str, AgentHandshake]:
        return dict(self._handshakes)

    def on_message(self, h: MessageHandler) -> None:
        self._msg_handlers.append(h)

    def on_approval_needed(self, h: ApprovalHandler) -> None:
        self._approval_handler = h

    async def start(self) -> None:
        self._config.resolve_paths()
        self._running = True
        self._listener = asyncio.create_task(self._listen(), name="p2p:listen")
        self._reconnector = asyncio.create_task(self._reconnect_loop(), name="p2p:reconnect")

        # Register with relay servers and start reading incoming messages
        self._relay_client = None
        if self._config.relay_addresses:
            from p2p.transport.relay import RelayClient
            self._relay_client = RelayClient(self._identity.peer_id, identity=self._identity)
            for addr in self._config.relay_addresses:
                try:
                    mux = await self._relay_client.connect_relay(
                        addr, on_new_peer=self._on_relay_peer,
                    )
                    logger.info("Registered with relay %s (reading)", addr)
                except Exception:
                    logger.warning("Failed to register with relay %s", addr, exc_info=True)

        logger.info("P2P node %s started on :%d", self._identity.peer_id, self._config.listen_port)

    async def _on_relay_peer(self, relay_conn) -> None:
        """Called by RelayMux when a new peer sends us data via relay.

        Performs Noise_XX handshake + identity binding over the relay
        connection, then treats it as a normal authenticated connection.
        """
        pid = relay_conn.peer_id
        logger.info("Incoming relay connection from %s", pid)

        try:
            from p2p.transport.websocket import (
                _noise_responder_raw,
                _x25519_key,
            )
            session, verified_pid = await asyncio.wait_for(
                _noise_responder_raw(relay_conn, self._identity),
                timeout=30,
            )
            # Wrap the relay connection with Noise encryption
            from p2p.transport.websocket import EncryptedRelayConnection
            enc_conn = EncryptedRelayConnection(relay_conn, session, verified_pid)
            await self._on_conn(enc_conn, initiator=False)
        except Exception:
            logger.warning("Relay Noise handshake failed with %s", pid, exc_info=True)

    async def stop(self) -> None:
        self._running = False
        for pid, conn in self._connections.items():
            try:
                env = self._envelope(pid, MessageType.STATUS,
                                     StatusPayload(status="offline_soon").model_dump())
                await conn.send(encode(env))
            except Exception:
                pass
        for t in list(self._tasks.values()):
            t.cancel()
        for t in [self._listener, self._reconnector]:
            if t:
                t.cancel()
        all_tasks = list(self._tasks.values()) + [t for t in [self._listener, self._reconnector] if t]
        for t in all_tasks:
            try:
                await t
            except asyncio.CancelledError:
                pass
        await self._transport.shutdown()
        self._connections.clear()
        self._tasks.clear()
        self._handshakes.clear()

    async def send(self, peer_id: str, msg_type: MessageType | str,
                   payload: dict, *, in_reply_to: str = "") -> str:
        if isinstance(msg_type, str):
            msg_type = MessageType(msg_type)
        conn = self._connections.get(peer_id)
        if not conn or not conn.is_open:
            raise PeerUnreachable(peer_id, "not connected")
        env = self._envelope(peer_id, msg_type, payload, in_reply_to)
        await conn.send(encode(env))
        return env.id

    async def connect_to_peer(self, peer_id: str) -> bool:
        if peer_id in self._connections and self._connections[peer_id].is_open:
            return True
        rec = self._peer_store.get(peer_id)

        # Try direct connection via stored addresses
        if rec and rec.addresses:
            for addr in rec.addresses:
                try:
                    conn = await self._transport.connect(addr)
                    await self._on_conn(conn, initiator=True)
                    return True
                except Exception:
                    continue

        # Try via relay — get a virtual connection and do Noise handshake
        if self._relay_client:
            for addr, mux in self._relay_client._muxes.items():
                if not mux.is_open:
                    continue
                try:
                    relay_conn = mux.get_connection(peer_id)
                    from p2p.transport.websocket import (
                        _noise_initiator_raw,
                        EncryptedRelayConnection,
                    )
                    session, verified_pid = await asyncio.wait_for(
                        _noise_initiator_raw(relay_conn, self._identity),
                        timeout=30,
                    )
                    enc_conn = EncryptedRelayConnection(relay_conn, session, verified_pid)
                    await self._on_conn(enc_conn, initiator=True)
                    logger.info("Connected to %s via relay %s", peer_id, addr)
                    return True
                except Exception:
                    logger.debug("Relay connect to %s via %s failed", peer_id, addr, exc_info=True)
                    continue

        return False

    async def approve_handshake(self, peer_id: str, trust_level: int = 1,
                                scope_instructions: str = "") -> bool:
        hs = self._handshakes.get(peer_id)
        if not hs:
            return False
        msgs = await hs.owner_decision(True, trust_level, scope_instructions)
        conn = self._connections.get(peer_id)
        if conn:
            for e in msgs:
                await conn.send(encode(e))
        if hs.is_complete:
            del self._handshakes[peer_id]
        return True

    async def reject_handshake(self, peer_id: str) -> bool:
        hs = self._handshakes.get(peer_id)
        if not hs:
            return False
        msgs = await hs.owner_decision(False)
        conn = self._connections.get(peer_id)
        if conn:
            for e in msgs:
                await conn.send(encode(e))
            await conn.close()
        self._handshakes.pop(peer_id, None)
        self._connections.pop(peer_id, None)
        return True

    async def disconnect_peer(self, peer_id: str, *, remove: bool = False) -> bool:
        """Disconnect from a peer. Optionally remove them from the peer store.

        Sends STATUS(offline_soon) before closing. If remove=True, also
        deletes the peer record so they won't auto-reconnect and would
        need a fresh handshake to connect again.
        """
        conn = self._connections.get(peer_id)
        if not conn:
            if remove:
                return self._peer_store.remove(peer_id)
            return False

        # Notify peer we're disconnecting
        try:
            env = self._envelope(peer_id, MessageType.STATUS,
                                 StatusPayload(status="offline_soon").model_dump())
            await conn.send(encode(env))
        except Exception:
            pass

        await conn.close()
        self._connections.pop(peer_id, None)

        # Cancel reader task
        task = self._tasks.pop(peer_id, None)
        if task:
            task.cancel()

        # Clean up any pending handshake
        self._handshakes.pop(peer_id, None)

        if remove:
            self._peer_store.remove(peer_id)

        logger.info("Disconnected from %s (remove=%s)", peer_id, remove)
        return True

    def health(self) -> dict:
        return {
            "peer_id": self._identity.peer_id,
            "running": self._running,
            "connected": {p: c.is_open for p, c in self._connections.items()},
            "handshakes": {p: h.state.value for p, h in self._handshakes.items()},
            "known_peers": len(self._peer_store),
        }

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

    def _envelope(self, to: str, t: MessageType, payload: dict,
                  reply: str = "") -> A2AEnvelope:
        env = A2AEnvelope(
            sender_id=self._identity.peer_id, recipient_id=to,
            type=t.value if isinstance(t, MessageType) else t,
            payload=payload, in_reply_to=reply,
        )
        env.signature_hex = self._identity.sign(env.signing_material()).hex()
        return env

    async def _listen(self) -> None:
        try:
            async for conn in self._transport.listen(
                self._config.listen_host, self._config.listen_port,
            ):
                if not self._running:
                    break
                asyncio.create_task(self._on_conn(conn, initiator=False))
        except asyncio.CancelledError:
            pass

    async def _on_conn(self, conn: Connection, *, initiator: bool) -> None:
        pid = conn.peer_id
        existing = self._peer_store.get(pid)
        if existing and existing.trust_level > 0:
            self._connections[pid] = conn
            self._peer_store.update_last_seen(pid)
            self._start_reader(pid, conn)
            try:
                env = self._envelope(pid, MessageType.STATUS,
                                     StatusPayload(status="online").model_dump())
                await conn.send(encode(env))
            except Exception:
                pass
            return

        self._connections[pid] = conn

        async def cb(p, h):
            if self._approval_handler:
                await self._approval_handler(p, h)

        hs = AgentHandshake(
            self._identity, self._peer_store, cb,
            our_hello=self._hello, is_initiator=initiator,
        )
        hs.set_remote_peer_id(pid)
        self._handshakes[pid] = hs
        for e in hs.start_messages():
            await conn.send(encode(e))
        self._start_reader(pid, conn)

    def _start_reader(self, pid: str, conn: Connection) -> None:
        old = self._tasks.pop(pid, None)
        if old:
            old.cancel()
        self._tasks[pid] = asyncio.create_task(self._reader(pid, conn))

    async def _reader(self, pid: str, conn: Connection) -> None:
        try:
            while conn.is_open and self._running:
                try:
                    raw = await conn.recv()
                except ConnectionError:
                    break
                try:
                    env = decode(raw)
                except Exception:
                    logger.debug("Malformed message from %s", pid)
                    continue

                # Dedup: reject replayed messages
                if env.id in self._seen_set:
                    logger.debug("Duplicate %s from %s", env.id[:8], pid)
                    continue
                if len(self._seen_queue) == self._seen_queue.maxlen:
                    self._seen_set.discard(self._seen_queue[0])
                self._seen_queue.append(env.id)
                self._seen_set.add(env.id)

                # Signature verification for non-handshake messages
                if env.type not in _HS_TYPES:
                    peer_rec = self._peer_store.get(pid)
                    if peer_rec and peer_rec.public_key_pem:
                        pub = PeerIdentity.load_public_key_pem(peer_rec.public_key_pem)
                        if not env.signature_hex or not PeerIdentity.verify(
                            env.signing_material(),
                            bytes.fromhex(env.signature_hex),
                            pub,
                        ):
                            logger.warning("Bad/missing signature from %s, dropping", pid)
                            continue

                await self._dispatch(pid, env)
        except asyncio.CancelledError:
            pass
        finally:
            self._connections.pop(pid, None)
            self._tasks.pop(pid, None)

    async def _dispatch(self, pid: str, env: A2AEnvelope) -> None:
        # Reject messages claiming to be from a different peer
        if env.sender_id and env.sender_id != pid:
            logger.warning("sender_id mismatch: %s vs conn %s — dropping", env.sender_id, pid)
            return

        if env.type in _HS_TYPES:
            hs = self._handshakes.get(pid)
            if not hs:
                return
            msgs = await hs.handle_message(env)
            conn = self._connections.get(pid)
            if conn:
                for e in msgs:
                    await conn.send(encode(e))
            if hs.is_complete:
                self._handshakes.pop(pid, None)
            elif hs.is_failed:
                self._handshakes.pop(pid, None)
                if conn:
                    await conn.close()
                self._connections.pop(pid, None)
            return

        if env.type == MessageType.PING.value:
            conn = self._connections.get(pid)
            if conn:
                await conn.send(encode(
                    self._envelope(pid, MessageType.PONG, {}, env.id)))
            return
        if env.type in (MessageType.PONG.value, MessageType.ACK.value):
            return
        if env.type == MessageType.STATUS.value:
            s = StatusPayload.model_validate(env.payload)
            if s.status == "offline_soon":
                self._peer_store.update_last_seen(pid)
            return

        # ACK application messages
        conn = self._connections.get(pid)
        if conn:
            try:
                await conn.send(encode(self._envelope(pid, MessageType.ACK, {}, env.id)))
            except Exception:
                pass

        for h in self._msg_handlers:
            try:
                await h(env)
            except Exception:
                logger.exception("Handler error for %s", pid)

    async def _reconnect_loop(self) -> None:
        try:
            while self._running:
                await asyncio.sleep(self._config.reconnect_interval_seconds)
                if not self._running:
                    break
                for rec in self._peer_store.list_all():
                    if rec.trust_level <= 0:
                        continue
                    if rec.peer_id in self._connections and self._connections[rec.peer_id].is_open:
                        continue
                    try:
                        await self.connect_to_peer(rec.peer_id)
                    except Exception:
                        pass
        except asyncio.CancelledError:
            pass
