"""
Agent-facing meta tools for searching and loading functions.

These tools are used by the agent to discover and load capabilities from the registry.
"""

import asyncio
from collections import defaultdict
from typing import Annotated, Any

from langchain_core.tools import tool as langchain_tool
from loguru import logger
from pydantic import Field
from universal_mcp.types import ToolFormat

from universal_mcp.agentr.registry import AgentrRegistry
from universal_mcp.agents.codeact0.prompts import build_tool_definitions


async def _generate_unconnected_links(
    tool_ids: list[str] | dict[str, list[str]], tool_registry: AgentrRegistry
) -> dict[str, str]:
    """For a given list of tool_ids, validates the tools and returns a list of links for the apps that have not been logged in"""
    apps = set()
    for tool_id in tool_ids:
        app_id = tool_id.split("__")[0]
        apps.add(app_id)
    logger.info(f"Unique apps: {apps}")
    connected = await tool_registry.list_connected_apps()
    connected_app_ids = {c["app_id"] for c in connected}
    logger.info(f"Connected apps: {connected_app_ids}")
    unconnected = apps - connected_app_ids
    logger.info(f"Unconnected apps: {unconnected}")
    links = {}
    for app_id in unconnected:
        if app_id == "llm":
            continue
        text = await tool_registry.authorise_app(app_id=app_id)
        start = text.find(":") + 1
        end = text.find(". R", start)
        url = text[start:end].strip()
        links[app_id] = url
    return links


async def get_valid_tools(tool_ids, tool_registry: AgentrRegistry):
    """Return valid tool ids, context, definition and unconnected links"""
    # Step 1: Load the tools, only valid tools are loaded that are present in registry
    await tool_registry.load_tools(tool_ids)
    exported_tools = await tool_registry.export_tools(tool_ids, ToolFormat.NATIVE)
    # Step 2: Build the informational string for the agent.
    tool_definitions, new_tools_context = build_tool_definitions(exported_tools)

    # Step 3: Generate links for unconnected apps
    valid_tools = list(new_tools_context.keys())

    if not valid_tools:
        response_string = "Error: None of the provided tool IDs could be validated or loaded."
        return response_string, {}, [], ""
    unconnected_links = await _generate_unconnected_links(valid_tools, tool_registry)
    return valid_tools, tool_definitions, new_tools_context, list(unconnected_links.values())


def create_meta_tools(tool_registry: AgentrRegistry | None = None) -> dict[str, Any]:  # noqa: PLR0915
    """Create the meta tools for searching and loading tools"""

    @langchain_tool
    async def search_functions(
        queries: Annotated[
            list[str],
            Field(description="A list of search queries."),
        ],
    ) -> str:
        """
        Discovers tools using a list of search queries, returning results using a vector-based search.
        Returns up to 5 functions for each application per search query.
        Use multiple queries in a single request to search for different applications, or different tasks in the same application.

        **Search Strategy Tips:**
        When using `search_functions`, craft effective search queries:
        - ✅ Good queries are specific but not too narrow: ["send email", "calendar events"]
        - ❌ Too narrow: ["send email using gmail API version 2 with OAuth2"]
        - ❌ Too broad: ["email", "data", "file"]

        If your first search doesn't find what you need:
        1. Try synonyms (e.g., "email" → "mail", "message" | "calendar" → "schedule", "events")
        2. Try broader terms (e.g., "gmail send" → "email send" | "google sheets read" → "spreadsheet read")
        3. Try related concepts (e.g., "send email" → "compose message" | "create meeting" → "schedule event")
        4. After 3 unsuccessful attempts, inform the user that the functionality may not be available
        """
        if not tool_registry:
            return "Tool registry not provided."

        if not queries:
            return "No search queries provided."

        # 2. Parallel Search across all queries
        # We fetch more results initially to ensure we can find up to 5 unique apps
        tasks = [tool_registry.search_tools(query=q, limit=40, distance_threshold=0.75) for q in queries]
        search_results = await asyncio.gather(*tasks)

        # 3. Connection Status and Grouping
        connections = await tool_registry.list_connected_apps()
        connected_app_ids = {c["app_id"] for c in connections}

        # Group tools by app: { app_id: { tool_id: tool_dict } }
        # Using a dict for tools ensures deduplication across multiple query results
        app_groups = defaultdict(dict)

        for tool_list in search_results:
            query_app_counts = defaultdict(int)
            for tool in tool_list:
                app_id = tool.get("app_id", "unknown")
                t_id = tool.get("id")

                if t_id and query_app_counts[app_id] < 5:
                    app_groups[app_id][t_id] = tool
                    query_app_counts[app_id] += 1

        # 4. Sort apps by connection status
        # We prioritize apps that already have connections if possible
        sorted_apps = sorted(app_groups.keys(), key=lambda x: x in connected_app_ids, reverse=True)
        selected_apps = sorted_apps

        if not selected_apps:
            return "No relevant functions were found."

        # 5. Formatting Output
        result_parts = []
        found_connected = set()

        for app_id in selected_apps:
            tools = app_groups[app_id].values()
            is_connected = app_id in connected_app_ids
            if is_connected:
                found_connected.add(app_id)

            status_str = "connected" if is_connected else "NOT connected"

            result_parts.append(f"Tools from {app_id} (status: {status_str} by user):")
            for tool in tools:
                desc = tool.get("description", "").split("Context:")[0].strip()
                result_parts.append(f" - {tool['id']}: {desc}")
            result_parts.append("")

        # 6. Status Summary & Footer
        if not found_connected and selected_apps:
            result_parts.append(
                "Connection Status: None of the apps found are connected. Please ask the user to choose an app to connect."
            )
        elif len(found_connected) > 1:
            result_parts.append(
                f"Connection Status: Multiple apps are connected ({', '.join(found_connected)}). Please ask the user to select one, unless they have already specified their choice or the apps do not seem relevant."
            )

        result_parts.append(
            "Call load_functions to select the required functions. For any unconnected apps, calling load_functions will provide a link to connect to the app for the user in addition to loading the functions."
        )

        return "\n".join(result_parts)

    @langchain_tool
    async def load_functions(tool_ids: list[str]) -> str:
        """
        Loads specified functions and returns their Python signatures and docstrings.
        This makes the functions available for use inside the 'execute_python_code' tool.
        The agent MUST use the returned information to understand how to call the functions correctly.

        Args:
            tool_ids: A list of function IDs in the format 'app__function'. Example: ['google_mail__send_email']

        Returns:
            A string containing the signatures and docstrings of the successfully loaded functions,
            ready for the agent to use in its code.
        """
        if not tool_registry:
            return "Tool registry not provided."
        if not tool_ids:
            return "No tool IDs provided to load."

        valid_tools, tool_definitions, new_tools_context, unconnected_links = await get_valid_tools(
            tool_ids, tool_registry
        )

        result_parts = [
            f"Successfully loaded {len(valid_tools)} functions. They are now available for use inside `execute_python_code`:",
            "\n".join(tool_definitions),
        ]

        response_string = "\n\n".join(result_parts)
        unconnected_links = "\n".join(unconnected_links)

        return response_string, new_tools_context, valid_tools, unconnected_links

    return {
        "search_functions": search_functions,
        "load_functions": load_functions,
    }
