"""Build the tool registry available to a run.

Combines built-in tools with the tools of every active connector installed in the
workspace. Connector tool handlers decrypt the credential only when the tool is
actually invoked, so plaintext secrets live in memory for the shortest time.
"""

from __future__ import annotations

import uuid

from sqlalchemy.orm import Session

from legate.agent.builtin_tools import BUILTIN_REGISTRY
from legate.agent.tooling import ToolContext, ToolRegistry, ToolResult, ToolSpec
from legate.connectors.base import Connector
from legate.connectors.registry import get_connector
from legate.observability.logging import get_logger
from legate.repositories import connectors as connectors_repo
from legate.repositories import credentials as credentials_repo
from legate.vault import VaultError, get_vault

logger = get_logger(__name__)


def _make_handler(connector: Connector, connector_id: uuid.UUID, tool_name: str):
    async def handler(args: dict, ctx: ToolContext) -> ToolResult:
        if ctx.db is None:
            return ToolResult(ok=False, error="no database session for connector call")
        row = connectors_repo.get(ctx.db, workspace_id=ctx.workspace_id, connector_id=connector_id)
        if row is None:
            return ToolResult(ok=False, error="connector not found")
        credential = credentials_repo.get_for_connector(ctx.db, connector_id=connector_id)
        if credential is None:
            return ToolResult(ok=False, error="connector has no stored credentials")
        try:
            secrets = get_vault().decrypt(
                credential.ciphertext, credential.wrapped_dek, credential.key_version
            )
        except VaultError as exc:
            return ToolResult(ok=False, error=f"could not decrypt credentials: {exc}")
        return await connector.execute(tool_name, args, row.config_json or {}, secrets, ctx)

    return handler


def mcp_specs_for_row(row) -> list[tuple[ToolSpec, str]]:
    """Turn a connector row's cached MCP tools into (registry spec, server tool name).

    MCP tools are namespaced ``mcp.<connector-id8>.<tool>`` so they are unique
    across servers and group under a single "mcp" key in the UI, while the server
    is still called with its original tool name.
    """
    out: list[tuple[ToolSpec, str]] = []
    for tool in (row.config_json or {}).get("mcp_tools", []):
        original = tool.get("tool")
        if not original:
            continue
        spec = ToolSpec(
            name=f"mcp.{str(row.id)[:8]}.{original}",
            description=tool.get("description", ""),
            parameters=tool.get("parameters") or {"type": "object", "properties": {}},
            is_sensitive=bool(tool.get("sensitive", True)),
            scopes=["mcp:call"],
        )
        out.append((spec, original))
    return out


def workspace_mcp_specs(db: Session, workspace_id: uuid.UUID) -> list[ToolSpec]:
    """Every cached MCP tool spec across the workspace's installed MCP servers."""
    specs: list[ToolSpec] = []
    for row in connectors_repo.list_for_workspace(db, workspace_id=workspace_id):
        if row.type == "mcp":
            specs.extend(spec for spec, _ in mcp_specs_for_row(row))
    return specs


def build_workspace_tool_registry(db: Session, workspace_id: uuid.UUID) -> ToolRegistry:
    registry = ToolRegistry()
    # Built-in tools first.
    for spec in BUILTIN_REGISTRY.specs():
        registered = BUILTIN_REGISTRY.get(spec.name)
        if registered is not None:
            registry.register(spec, registered.handler)

    # Then tools from each active connector installed in this workspace.
    for row in connectors_repo.list_for_workspace(db, workspace_id=workspace_id):
        connector = get_connector(row.type)
        if connector is None:
            continue

        # MCP connectors expose dynamic, cached tools (namespaced), each calling
        # the server with its original tool name.
        if row.type == "mcp":
            for spec, original in mcp_specs_for_row(row):
                if registry.get(spec.name) is not None:
                    continue
                registry.register(spec, _make_handler(connector, row.id, original))
            continue

        for spec in connector.list_tools():
            if registry.get(spec.name) is not None:
                logger.warning(
                    "tool %s already registered; skipping connector %s", spec.name, row.id
                )
                continue
            registry.register(spec, _make_handler(connector, row.id, spec.name))
    return registry
