"""A2ANode — the main P2P runtime.

Manages connections, routing, handshakes, and message delivery.
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Awaitable, Callable

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

logger = logging.getLogger(__name__)

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


class A2ANode:
    """The P2P node runtime.

    Usage::

        node = A2ANode(identity, peer_store, config)
        node.on_message(my_handler)
        node.on_approval_needed(my_approval_handler)
        await node.start()
        await node.send(peer_id, MessageType.TEXT, {"text": "hello"})
        ...
        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

        # Our HELLO payload
        self._our_hello = HelloPayload(
            agent_name=agent_name or "A2AAgent",
            owner_display_name=owner_display_name,
            capabilities=capabilities or [],
            introduction=introduction,
        )

        # Transport
        self._transport = WSTransport(identity)

        # Active connections (peer_id → Connection)
        self._connections: dict[str, Connection] = {}

        # Pending handshakes (peer_id → AgentHandshake)
        self._pending_handshakes: dict[str, AgentHandshake] = {}

        # Callbacks
        self._message_handlers: list[MessageHandler] = []
        self._approval_handler: ApprovalHandler | None = None

        # Tasks
        self._listener_task: asyncio.Task | None = None
        self._reader_tasks: dict[str, asyncio.Task] = {}
        self._reconnect_task: asyncio.Task | None = None
        self._running = False

    @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._pending_handshakes)

    def on_message(self, handler: MessageHandler) -> None:
        """Register a handler for incoming (non-handshake) messages."""
        self._message_handlers.append(handler)

    def on_approval_needed(self, handler: ApprovalHandler) -> None:
        """Register handler for when a new peer needs owner approval."""
        self._approval_handler = handler

    async def start(self) -> None:
        """Start the node: listen for connections + connect to known peers."""
        self._config.resolve_paths()
        self._running = True

        # Start listener
        self._listener_task = asyncio.create_task(
            self._listen_loop(),
            name="a2a:listener",
        )

        # Start reconnection loop for known peers
        self._reconnect_task = asyncio.create_task(
            self._reconnect_loop(),
            name="a2a:reconnect",
        )

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

    async def stop(self) -> None:
        """Gracefully stop the node."""
        self._running = False

        # Notify peers we're going offline
        for peer_id, conn in self._connections.items():
            try:
                status_env = self._make_envelope(
                    peer_id, MessageType.STATUS,
                    StatusPayload(status="offline_soon").model_dump(),
                )
                await conn.send(encode(status_env))
            except Exception:
                pass

        # Cancel tasks
        for task in self._reader_tasks.values():
            task.cancel()
        if self._listener_task:
            self._listener_task.cancel()
        if self._reconnect_task:
            self._reconnect_task.cancel()

        # Wait for cancellation
        tasks = list(self._reader_tasks.values())
        if self._listener_task:
            tasks.append(self._listener_task)
        if self._reconnect_task:
            tasks.append(self._reconnect_task)
        for task in tasks:
            try:
                await task
            except asyncio.CancelledError:
                pass

        # Close transport
        await self._transport.shutdown()

        self._connections.clear()
        self._reader_tasks.clear()
        self._pending_handshakes.clear()

        logger.info("A2A node stopped")

    async def send(
        self,
        peer_id: str,
        msg_type: MessageType | str,
        payload: dict,
        *,
        in_reply_to: str = "",
    ) -> str:
        """Send a message to a peer. Returns the message ID.

        Raises PeerUnreachable if the peer is not connected.
        """
        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")

        envelope = self._make_envelope(peer_id, msg_type, payload, in_reply_to)
        await conn.send(encode(envelope))
        return envelope.id

    async def connect_to_peer(self, peer_id: str) -> bool:
        """Attempt to connect to a known peer.

        Returns True if connection + handshake initiated successfully.
        """
        if peer_id in self._connections and self._connections[peer_id].is_open:
            return True  # Already connected

        record = self._peer_store.get(peer_id)
        if not record or not record.addresses:
            logger.warning("No addresses for peer %s", peer_id)
            return False

        for addr in record.addresses:
            try:
                conn = await self._transport.connect(addr)
                await self._on_new_connection(conn, is_initiator=True)
                return True
            except Exception:
                logger.debug(
                    "Failed to connect to %s at %s", 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:
        """Approve a pending handshake for a peer. Returns True if handled."""
        hs = self._pending_handshakes.get(peer_id)
        if not hs:
            logger.warning("No pending handshake for %s", peer_id)
            return False

        responses = await hs.owner_decision(
            approved=True,
            trust_level=trust_level,
            scope_instructions=scope_instructions,
        )

        conn = self._connections.get(peer_id)
        if conn:
            for env in responses:
                await conn.send(encode(env))

        if hs.is_complete:
            del self._pending_handshakes[peer_id]
            logger.info("Handshake approved and completed for %s", peer_id)

        return True

    async def reject_handshake(self, peer_id: str) -> bool:
        """Reject a pending handshake. Returns True if handled."""
        hs = self._pending_handshakes.get(peer_id)
        if not hs:
            return False

        responses = await hs.owner_decision(approved=False)

        conn = self._connections.get(peer_id)
        if conn:
            for env in responses:
                await conn.send(encode(env))
            await conn.close()

        del self._pending_handshakes[peer_id]
        self._connections.pop(peer_id, None)
        return True

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

    def _make_envelope(
        self,
        recipient_id: str,
        msg_type: MessageType,
        payload: dict,
        in_reply_to: str = "",
    ) -> A2AEnvelope:
        env = A2AEnvelope(
            sender_id=self._identity.peer_id,
            recipient_id=recipient_id,
            type=msg_type.value if isinstance(msg_type, MessageType) else msg_type,
            payload=payload,
            in_reply_to=in_reply_to,
        )
        # Sign
        sig = self._identity.sign(env.signing_material())
        env.signature_hex = sig.hex()
        return env

    async def _listen_loop(self) -> None:
        """Accept incoming connections from the transport."""
        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_new_connection(conn, is_initiator=False),
                    name=f"a2a:conn:{conn.peer_id[:12]}",
                )
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("Listener loop crashed")

    async def _on_new_connection(
        self, conn: Connection, *, is_initiator: bool,
    ) -> None:
        """Handle a new connection (either inbound or outbound)."""
        peer_id = conn.peer_id
        logger.info("New connection with %s (initiator=%s)", peer_id, is_initiator)

        # Check if this is a known, trusted peer (skip handshake)
        existing = self._peer_store.get(peer_id)
        if existing and existing.trust_level > 0:
            # Known peer — no handshake needed, just reconnect
            self._connections[peer_id] = conn
            self._peer_store.update_last_seen(peer_id)
            self._start_reader(peer_id, conn)

            # Send STATUS online
            try:
                status_env = self._make_envelope(
                    peer_id, MessageType.STATUS,
                    StatusPayload(status="online").model_dump(),
                )
                await conn.send(encode(status_env))
            except Exception:
                pass

            logger.info("Reconnected with known peer %s (%s)", peer_id, existing.alias)
            return

        # New peer — start handshake
        self._connections[peer_id] = conn

        async def approval_callback(pid: str, hello_dict: dict) -> None:
            if self._approval_handler:
                await self._approval_handler(pid, hello_dict)

        hs = AgentHandshake(
            identity=self._identity,
            peer_store=self._peer_store,
            on_approval_needed=approval_callback,
            our_hello=self._our_hello,
            is_initiator=is_initiator,
        )
        hs.set_remote_peer_id(peer_id)
        self._pending_handshakes[peer_id] = hs

        # Send initial handshake messages
        for env in hs.start_messages():
            await conn.send(encode(env))

        # Start reading from connection
        self._start_reader(peer_id, conn)

    def _start_reader(self, peer_id: str, conn: Connection) -> None:
        """Start an asyncio task to read messages from a connection."""
        # Cancel existing reader if any
        old = self._reader_tasks.pop(peer_id, None)
        if old:
            old.cancel()

        task = asyncio.create_task(
            self._reader_loop(peer_id, conn),
            name=f"a2a:reader:{peer_id[:12]}",
        )
        self._reader_tasks[peer_id] = task

    async def _reader_loop(self, peer_id: str, conn: Connection) -> None:
        """Read messages from a connection and dispatch them."""
        try:
            while conn.is_open and self._running:
                try:
                    raw = await conn.recv()
                except ConnectionError:
                    break

                try:
                    envelope = decode(raw)
                except Exception:
                    logger.warning("Bad message from %s", peer_id, exc_info=True)
                    continue

                await self._dispatch(peer_id, envelope)

        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("Reader loop crashed for %s", peer_id)
        finally:
            # Clean up
            self._connections.pop(peer_id, None)
            self._reader_tasks.pop(peer_id, None)
            logger.info("Connection closed with %s", peer_id)

    async def _dispatch(self, peer_id: str, envelope: A2AEnvelope) -> None:
        """Route a message to the handshake FSM or application handlers."""
        msg_type = envelope.type

        # Handshake messages
        handshake_types = {
            MessageType.HELLO.value,
            MessageType.CHALLENGE.value,
            MessageType.CHALLENGE_RESPONSE.value,
            MessageType.APPROVAL_REQUEST.value,
            MessageType.APPROVAL_GRANT.value,
            MessageType.APPROVAL_DENY.value,
        }

        if msg_type in handshake_types:
            hs = self._pending_handshakes.get(peer_id)
            if not hs:
                logger.warning("Handshake message from %s but no pending handshake", peer_id)
                return

            responses = await hs.handle_message(envelope)
            conn = self._connections.get(peer_id)
            if conn:
                for env in responses:
                    await conn.send(encode(env))

            # Check if handshake completed
            if hs.is_complete:
                del self._pending_handshakes[peer_id]
                logger.info("Handshake completed with %s", peer_id)
            elif hs.is_failed:
                del self._pending_handshakes[peer_id]
                if conn:
                    await conn.close()
                self._connections.pop(peer_id, None)
                logger.info("Handshake failed/rejected with %s", peer_id)

            return

        # ACK handling (simple)
        if msg_type == MessageType.ACK.value:
            return

        # PING/PONG
        if msg_type == MessageType.PING.value:
            pong = self._make_envelope(
                peer_id, MessageType.PONG, {},
                in_reply_to=envelope.id,
            )
            conn = self._connections.get(peer_id)
            if conn:
                await conn.send(encode(pong))
            return

        if msg_type == MessageType.PONG.value:
            return

        # STATUS
        if msg_type == MessageType.STATUS.value:
            status = StatusPayload.model_validate(envelope.payload)
            logger.info("Peer %s status: %s", peer_id, status.status)
            if status.status == "offline_soon":
                self._peer_store.update_last_seen(peer_id)
            return

        # Application messages — dispatch to handlers
        # Send ACK first
        ack = self._make_envelope(
            peer_id, MessageType.ACK, {},
            in_reply_to=envelope.id,
        )
        conn = self._connections.get(peer_id)
        if conn:
            try:
                await conn.send(encode(ack))
            except Exception:
                pass

        # Dispatch to handlers
        for handler in self._message_handlers:
            try:
                await handler(envelope)
            except Exception:
                logger.exception("Message handler error for %s", peer_id)

    async def _reconnect_loop(self) -> None:
        """Periodically try to reconnect to known peers that are offline."""
        try:
            while self._running:
                await asyncio.sleep(self._config.reconnect_interval_seconds)
                if not self._running:
                    break

                for record in self._peer_store.list_all():
                    if record.trust_level <= 0:
                        continue  # Don't auto-reconnect to unverified peers
                    if record.peer_id in self._connections:
                        if self._connections[record.peer_id].is_open:
                            continue  # Already connected

                    # Try to connect
                    try:
                        await self.connect_to_peer(record.peer_id)
                    except Exception:
                        pass

        except asyncio.CancelledError:
            pass

    def health(self) -> dict:
        """Return node health status."""
        return {
            "peer_id": self._identity.peer_id,
            "running": self._running,
            "connected_peers": {
                pid: conn.is_open for pid, conn in self._connections.items()
            },
            "pending_handshakes": {
                pid: hs.state.value
                for pid, hs in self._pending_handshakes.items()
            },
            "known_peers": len(self._peer_store),
        }
