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:
        """Search for agents on the MRP network by capability, prefix, or name.

        Use this to find agents that can help with a task. At least one
        search parameter is required. Results are NOT automatically saved
        to contacts — use mrp_add_contact to save an agent.

        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.
        """
        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,
        )

        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("Use mrp_add_contact to save an agent to your contacts.")
        return "\n".join(lines)

    @mcp.tool()
    async def mrp_add_contact(
        name: str,
        alias: str | None = None,
    ) -> str:
        """Add an agent to your contacts so you can send them messages.

        Look up an agent by name or public key and save them to your
        contacts list. Use mrp_discover first to find agents on the network.

        Args:
            name: The agent's display name or public key (from mrp_discover results).
            alias: Optional short alias for easier reference (e.g. "bob").

        Returns:
            Confirmation with the agent's info, or an error if not found.
        """
        # Try searching by name first, then treat as public key
        results = await asyncio.to_thread(
            agent.discover_search,
            capabilities=None,
            capability_prefix=None,
            name=name,
        )

        # Find exact or best match
        match = None
        for a in results:
            if a.public_key == name:
                match = a
                break
            if a.display_name and a.display_name.lower() == name.lower():
                match = a
                break
        if not match and results:
            match = results[0]

        if not match:
            return f"No agent found matching '{name}'. Use mrp_discover to search first."

        contact = state.add_contact(match, alias=alias)
        display = contact.display_name or "(unnamed)"
        alias_str = f" (alias: {alias})" if alias else ""
        caps = ", ".join(contact.capabilities) if contact.capabilities else "none"
        return (
            f"Added **{display}**{alias_str} to contacts.\n"
            f"Capabilities: {caps}\n"
            f"You can now use their name in mrp_send."
        )

    @mcp.tool()
    async def mrp_list_contacts() -> str:
        """List all agents you've added to your contacts.

        Shows agents saved via mrp_add_contact, along with their
        capabilities and any aliases you've set.

        Returns:
            A list of your contacts with their names and capabilities.
        """
        if not state.contacts:
            return "No contacts yet. Use mrp_discover to find agents, then mrp_add_contact to save them."

        lines = ["Your 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)
