from __future__ import annotations

import asyncio

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

from mrp_mcp.state import SessionState


def register(mcp: FastMCP, agent: Agent, state: SessionState) -> None:
    @mcp.tool()
    async def mrp_discover(
        capability: str | None = None,
        capability_prefix: str | None = None,
        name: str | None = None,
    ) -> str:
        """Find agents on the MRP network by capability, prefix, or name.

        Use this to discover agents that can help with a task. At least one
        search parameter is required.

        Search modes (can be combined):
        - capability: exact match (e.g. "text:translate")
        - capability_prefix: namespace match (e.g. "text:" finds all text agents)
        - name: substring match on agent display names

        Capabilities use colon-separated namespaces like:
        - text:translate, text:summarize, text:chat
        - code:review, code:generate, code:explain
        - image:generate, image:analyze

        Args:
            capability: Exact capability to search for (e.g. "text:translate")
            capability_prefix: Capability prefix to match (e.g. "text:")
            name: Agent display name substring to search for

        Returns:
            A list of agents with their names, capabilities, and when they were last active.
            These agents are saved to your contacts for easy reference in mrp_send.
        """
        if not capability and not capability_prefix and not name:
            return "Error: at least one of capability, capability_prefix, or name is required."

        caps = [capability] if capability else None
        results = await asyncio.to_thread(
            agent.discover_search,
            capabilities=caps,
            capability_prefix=capability_prefix,
            name=name,
        )

        for a in results:
            state.add_contact(a)

        if not results:
            parts = []
            if capability:
                parts.append(f"capability '{capability}'")
            if capability_prefix:
                parts.append(f"prefix '{capability_prefix}'")
            if name:
                parts.append(f"name '{name}'")
            return f"No agents found matching {', '.join(parts)}."

        lines = [f"Found {len(results)} agent(s):\n"]
        for i, a in enumerate(results, 1):
            display = a.display_name or "(unnamed)"
            lines.append(f"{i}. **{display}**")
            lines.append(f"   Public key: {a.public_key}")
            lines.append(f"   Capabilities: {', '.join(a.capabilities)}")
            lines.append(f"   Last active: {a.last_active_at.isoformat()}")
            lines.append("")

        lines.append("These agents are now in your contacts. Use their name in mrp_send.")
        return "\n".join(lines)

    @mcp.tool()
    async def mrp_list_contacts() -> str:
        """List all agents you've discovered in this session.

        Shows agents found via previous mrp_discover calls, along with
        their capabilities and any aliases you've set.

        Returns:
            A list of known agents with their names and capabilities.
        """
        if not state.contacts:
            return "No contacts yet. Use mrp_discover to find agents."

        lines = ["Known contacts:\n"]
        for pk, contact in state.contacts.items():
            display = contact.display_name or "(unnamed)"
            alias = f" (alias: {contact.alias})" if contact.alias else ""
            caps = ", ".join(contact.capabilities) if contact.capabilities else "none"
            lines.append(f"- **{display}**{alias}: {caps}")
        return "\n".join(lines)

    @mcp.tool()
    async def mrp_capabilities() -> str:
        """List all capabilities available on the MRP network.

        Use this to browse what capabilities exist before discovering specific agents.
        Shows each capability tag and how many agents offer it.

        Returns:
            A list of capability tags with agent counts, sorted by popularity.
        """
        result = await asyncio.to_thread(agent.list_capabilities)

        if not result.capabilities:
            return "No capabilities registered on the network."

        lines = [f"Found {len(result.capabilities)} capability tag(s):\n"]
        for c in result.capabilities:
            lines.append(f"- **{c.tag}** ({c.agent_count} agent{'s' if c.agent_count != 1 else ''})")
        return "\n".join(lines)
