import asyncio
import json

from .decider import decide_tools
from .errors import MCPError
from .mcp import MCPManager
from .utils import fetch_connectors


_ASK_FOR_TOOL = {
    "type": "function",
    "name": "ask_for_tool",
    "description": (
        "Request additional tools when you need capabilities not in your "
        "current tool set. Describe what you need."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "description": {
                "type": "string",
                "description": "What tools or capabilities you need",
            }
        },
        "required": ["description"],
    },
}


class Akashic:
    """Pre-flight MCP tool provider.

    Connects to MCP servers, handles OAuth, and exposes tools
    in OpenAI function-call format.  Does **not** wrap the LLM call —
    the developer keeps full control of their OpenAI client.

    Usage::

        def on_auth(urls):
            for u in urls:
                print(f"Sign in to {u['serverLabel']}: {u['authorizationUrl']}")

        akashic = Akashic(api_key="ak_xxx", on_auth=on_auth)
        client  = OpenAI()

        akashic.register("Push this to my repo")

        response = client.responses.create(
            model="gpt-5", input=data, tools=akashic.tools
        )
        result = akashic.call_tool("push_files", {"files": [...]})
    """

    def __init__(self, api_key: str, *, on_auth=None):
        self._api_key = api_key
        self._mcp = MCPManager()
        self._connectors: list[dict] | None = None
        self._on_auth = on_auth

    def _handle_auth(self, auth_urls: list[dict]) -> None:
        """Surface auth URLs via the callback, then block until OAuth completes."""
        if not auth_urls:
            return
        if self._on_auth:
            self._on_auth(auth_urls)
        asyncio.run(self._mcp.complete_auth(self._api_key))

    def register(self, data: str = "") -> None:
        """Pre-flight: connect to MCP servers, initiate OAuth if needed.

        When *data* is provided, a lightweight GPT-5-nano call selects only the
        tools relevant to the user's message, then connects only the MCP
        servers that own those tools.  When *data* is empty, all configured
        connectors are connected (backward-compatible).

        After this call, filtered tools are available via ``self.tools``.

        Args:
            data: The user's input for tool filtering.
        """
        if self._connectors is None:
            self._connectors = fetch_connectors(self._api_key)

        # Reset tool registrations so new filters can apply cleanly
        self._mcp.reset_tools()

        # No user input → connect all connectors (backward-compatible)
        if not data:
            _, auth_urls = asyncio.run(self._mcp.connect(self._connectors, self._api_key))
            self._handle_auth(auth_urls)
            return

        # Build flat list of {name, description} across all connectors
        all_tools: list[dict] = []
        tool_to_connectors: dict[str, list[int]] = {}  # tool name → connector indices
        for idx, conn in enumerate(self._connectors):
            for tool in conn.get("tools", []):
                all_tools.append(tool)
                tool_to_connectors.setdefault(tool["name"], []).append(idx)

        # If no tool metadata available (all connectors lack catalog info),
        # fall back to connecting everything
        if not all_tools:
            _, auth_urls = asyncio.run(self._mcp.connect(self._connectors, self._api_key))
            self._handle_auth(auth_urls)
            return

        # Ask GPT-5-nano which tools are relevant
        chosen_names = decide_tools(data, all_tools)

        # No tools needed → skip MCP connections
        if not chosen_names:
            return

        # Filter connectors to only those owning chosen tools
        needed_indices: set[int] = set()
        for name in chosen_names:
            for idx in tool_to_connectors.get(name, []):
                needed_indices.add(idx)

        filtered: list[dict] = []
        for idx in sorted(needed_indices):
            conn = self._connectors[idx]
            # Narrow allowedTools to just the chosen subset for this connector
            connector_tool_names = {t["name"] for t in conn.get("tools", [])}
            narrowed = [n for n in chosen_names if n in connector_tool_names]
            filtered.append({**conn, "allowedTools": narrowed})

        # Also include connectors not in the catalog (empty tools array)
        # since we can't filter what we don't know about
        for idx, conn in enumerate(self._connectors):
            if idx not in needed_indices and not conn.get("tools"):
                filtered.append(conn)

        _, auth_urls = asyncio.run(self._mcp.connect(filtered, self._api_key))
        self._handle_auth(auth_urls)

    @property
    def tools(self) -> list[dict]:
        """MCP tools in OpenAI function-call format, ready to pass to the LLM."""
        return self._mcp.tools + [_ASK_FOR_TOOL]

    def _expand_tools(self, description: str) -> str:
        """Load additional tools matching *description* from unloaded connectors."""
        loaded = set(self._mcp._tool_map)
        unloaded: list[dict] = []
        tool_to_connectors: dict[str, list[int]] = {}
        for idx, conn in enumerate(self._connectors or []):
            for tool in conn.get("tools", []):
                if tool["name"] not in loaded:
                    unloaded.append(tool)
                    tool_to_connectors.setdefault(tool["name"], []).append(idx)

        if not unloaded:
            return "No additional tools available. All configured tools are already loaded."

        chosen_names = decide_tools(description, unloaded)
        if not chosen_names:
            return f"No matching tools found for: {description}"

        needed_indices: set[int] = set()
        for name in chosen_names:
            for idx in tool_to_connectors.get(name, []):
                needed_indices.add(idx)

        filtered: list[dict] = []
        for idx in sorted(needed_indices):
            conn = self._connectors[idx]
            connector_tool_names = {t["name"] for t in conn.get("tools", [])}
            narrowed = [n for n in chosen_names if n in connector_tool_names]
            filtered.append({**conn, "allowedTools": narrowed})

        _, auth_urls = asyncio.run(self._mcp.connect(filtered, self._api_key))
        self._handle_auth(auth_urls)

        added = [n for n in chosen_names if n in self._mcp._tool_map]
        if added:
            return "Tools loaded: " + ", ".join(added)
        return "Failed to load requested tools. The matching connectors may require OAuth or be unreachable."

    def execute(self, output: list) -> list[dict]:
        """Execute Akashic tool calls from an LLM response.

        Filters *output* for ``function_call`` items whose tool name is
        known to Akashic, executes them via MCP, and returns a list of
        ``function_call_output`` dicts ready to append to the conversation.

        Tool calls for names **not** registered with Akashic are silently
        skipped so users can handle their own tools separately.

        Args:
            output: ``response.output`` from the Responses API — a list of
                output items (SDK objects or plain dicts).

        Returns:
            List of ``{"type": "function_call_output", "call_id": ...,
            "output": ...}`` dicts.
        """
        known_tools = {t["name"] for t in self.tools}
        results: list[dict] = []

        for item in output:
            # Support both SDK objects (attribute access) and plain dicts
            item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
            if item_type != "function_call":
                continue

            name = item.get("name") if isinstance(item, dict) else getattr(item, "name", None)
            call_id = item.get("call_id") if isinstance(item, dict) else getattr(item, "call_id", None)
            raw_args = item.get("arguments") if isinstance(item, dict) else getattr(item, "arguments", "{}")

            if name == "ask_for_tool":
                arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
                tool_result = self._expand_tools(arguments.get("description", ""))
                results.append({
                    "type": "function_call_output",
                    "call_id": call_id,
                    "output": tool_result,
                })
                continue

            if name not in known_tools:
                continue

            try:
                arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
            except json.JSONDecodeError as exc:
                raise MCPError(
                    f"Tool '{name}' received malformed arguments from the model. "
                    f"The raw arguments were not valid JSON: {exc}"
                ) from exc
            tool_result = self.call_tool(name, arguments)
            output_str = json.dumps(tool_result) if not isinstance(tool_result, str) else tool_result

            results.append({
                "type": "function_call_output",
                "call_id": call_id,
                "output": output_str,
            })

        return results

    def run(self, response, *, create, max_turns: int = 10) -> object:
        """Execute tool calls in a loop until the model stops calling tools.

        Takes an initial LLM response, executes any Akashic tool calls,
        appends the results to the conversation, re-sends to the model,
        and repeats until no more tool calls are made (or *max_turns* is hit).

        Args:
            response: The initial response object from ``client.responses.create()``.
            create: A callable that re-invokes the model, e.g.
                ``lambda msgs: client.responses.create(model=..., input=msgs, tools=...)``
            max_turns: Safety cap on tool-call rounds (default 10).

        Returns:
            The final response object (no more tool calls).
        """
        for _ in range(max_turns):
            tool_outputs = self.execute(response.output)
            if not tool_outputs:
                break
            follow_up = list(response.output) + tool_outputs
            response = create(follow_up)
        return response

    def call_tool(self, name: str, arguments: dict) -> dict:
        """Execute an MCP tool call by name, routing to the correct server."""
        return asyncio.run(self._mcp.call_tool(name, arguments))
