from __future__ import annotations

import asyncio
import json

from mcp.server.fastmcp import FastMCP
from mrp import Agent
from mrp.models import DiscoveredAgent

from mrp_mcp.formatting import format_message
from mrp_mcp.state import SessionState


def register(mcp: FastMCP, agent: Agent, state: SessionState) -> None:
    @mcp.tool()
    async def mrp_send(
        to: str,
        message: str,
        thread_id: str | None = None,
    ) -> str:
        """Send a message to another agent on the MRP network.

        Args:
            to: The recipient - can be an agent's display name, alias, or public key.
                Use mrp_discover first to find agents and add them to contacts.
            message: The message text to send.
            thread_id: Optional thread ID to organize this into a conversation.
                Use a descriptive ID like "translation-task-1".

        Returns:
            Confirmation with the message ID and delivery status.
        """
        recipient_pk = state.resolve_recipient(to)
        body = {"text": message}

        result = await asyncio.to_thread(
            agent.send, to=recipient_pk, body=body, thread=thread_id
        )

        contact = state.contacts.get(recipient_pk)
        name = contact.display_name if contact else recipient_pk[:12] + "..."
        return (
            f"Message sent to **{name}**.\n"
            f"Message ID: {result.message_id}\n"
            f"Thread: {thread_id or '(none)'}\n"
            f"Expires: {result.expires_at.isoformat()}"
        )

    @mcp.tool()
    async def mrp_reply(
        message_id: str,
        message: str,
    ) -> str:
        """Reply to a specific message you received.

        This automatically sets the recipient to the original sender and
        links the reply to the original message's conversation thread.

        Args:
            message_id: The ID of the message to reply to (from mrp_check_messages).
            message: The reply text.

        Returns:
            Confirmation with the reply's message ID.
        """
        original = await asyncio.to_thread(agent.client.get_message, message_id)
        body = {"text": message}
        result = await asyncio.to_thread(agent.reply, original, body)

        sender = state.contacts.get(original.sender_key)
        name = sender.display_name if sender else original.sender_key[:12] + "..."
        return (
            f"Reply sent to **{name}**.\n"
            f"Message ID: {result.message_id}\n"
            f"In reply to: {message_id}"
        )

    @mcp.tool()
    async def mrp_check_messages(
        from_agent: str | None = None,
        limit: int = 10,
    ) -> str:
        """Check for new incoming messages.

        Call this to see if any agents have sent you messages. Messages are
        marked as delivered once you read them.

        Args:
            from_agent: Optional filter - only show messages from this agent
                (name, alias, or public key).
            limit: Maximum number of messages to return (1-50, default 10).

        Returns:
            A list of new messages with sender, content, and message IDs
            (needed for mrp_reply).
        """
        sender_key = None
        if from_agent:
            sender_key = state.resolve_recipient(from_agent)

        # Check background poller's unread queue first
        messages = []
        if state.unread:
            for msg in list(state.unread):
                if sender_key and msg.sender_key != sender_key:
                    continue
                messages.append(msg)
                state.unread.remove(msg)
                if len(messages) >= limit:
                    break

        # If we need more, poll the relay
        if len(messages) < limit:
            result = await asyncio.to_thread(
                agent.client.poll_messages,
                status="sent",
                limit=limit - len(messages),
                sender_key=sender_key,
            )
            messages.extend(result.messages)

        # Add senders to contacts
        for msg in messages:
            if msg.sender_key not in state.contacts:
                try:
                    info = await asyncio.to_thread(
                        agent.client.get_agent, msg.sender_key
                    )
                    state.add_contact(
                        DiscoveredAgent(
                            public_key=info.public_key,
                            display_name=info.display_name,
                            last_active_at=info.last_active_at,
                            capabilities=info.capabilities,
                        )
                    )
                except Exception:
                    pass

        if not messages:
            return "No new messages."

        lines = [f"You have {len(messages)} new message(s):\n"]
        for msg in messages:
            lines.append("---")
            lines.append(format_message(msg, state))
            lines.append("")

        return "\n".join(lines)
