"""Tool & Integration Registry — single source of truth.

Tools and integrations self-register. The system prompt, tool definitions,
and contract enforcement are all built dynamically from this registry.

Adding a new integration:
    registry.register_integration("gcal", lambda: gcal.is_connected, "Google Calendar")
    registry.register_tool(
        name="list_events",
        description="List upcoming calendar events",
        handler=gcal_handler,
        integration="gcal",
        contract_action="read_calendar",
        input_schema={...},
    )

That's it. The agent automatically knows about it.
"""

from __future__ import annotations

import logging
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

log = logging.getLogger(__name__)

_CACHE_TTL = 60  # seconds — integration status cache lifetime


@dataclass
class ToolDef:
    """A registered tool."""

    name: str
    description: str
    input_schema: dict
    handler: Callable[..., Any]
    async_handler: Callable[..., Any] | None = None  # async version for direct await
    integration: str | None = None  # which integration this belongs to
    contract_action: str | None = None  # action name for contract check
    category: str = ""  # grouping for prompt display
    namespace: str = "misc"  # semantic namespace for filtering; see `agent/namespaces.py`


@dataclass
class IntegrationDef:
    """A registered integration."""

    name: str
    display_name: str
    description: str
    status_fn: Callable[[], bool]  # returns True if connected
    tools: list[str] = field(default_factory=list)


class ToolRegistry:
    """Central registry for tools and integrations.

    All tools register here. The system prompt, LLM tool definitions,
    and contract enforcement are built from this registry.
    """

    def __init__(self) -> None:
        self._tools: dict[str, ToolDef] = {}
        self._integrations: dict[str, IntegrationDef] = {}
        self._contract: Any = None  # set later to avoid circular imports
        self._mcp_client: Any = None  # MCPClientManager, set post-init
        self._namespace_resolver: Any = None  # NamespaceResolver, set post-init
        # Caches — avoid calling status_fn() on every message
        self._tool_defs_cache: list[dict] | None = None
        self._tool_defs_cache_time: float = 0
        self._connected_cache: list[str] | None = None
        self._connected_cache_time: float = 0
        self._capabilities_cache: str | None = None
        self._capabilities_cache_time: float = 0

    def set_contract(self, contract: Any) -> None:
        """Set the contract enforcer for tool execution."""
        self._contract = contract

    def set_namespace_resolver(self, resolver: Any) -> None:
        """Attach the `NamespaceResolver` so tool registrations can auto-resolve
        their namespace from the YAML map. Call this BEFORE populating the
        registry so tools pick up the right namespace at registration time.
        """
        self._namespace_resolver = resolver
        self._tool_defs_cache = None

    # ── Registration ──

    def register_tool(
        self,
        name: str,
        description: str,
        input_schema: dict,
        handler: Callable[..., Any],
        integration: str | None = None,
        contract_action: str | None = None,
        category: str = "",
        async_handler: Callable[..., Any] | None = None,
        namespace: str | None = None,
    ) -> None:
        """Register a tool. If integration is specified, links to it.

        ``namespace`` follows a 3-step resolution: (1) explicit kwarg wins,
        (2) otherwise look up the tool in the `NamespaceResolver` YAML map,
        (3) otherwise fall back to ``misc``. Callers can always override by
        passing an explicit namespace — useful for plugin-registered tools.
        """
        resolved_ns = namespace
        if not resolved_ns and self._namespace_resolver is not None:
            try:
                resolved_ns = self._namespace_resolver.namespace_for(name)
            except Exception:
                resolved_ns = None
        self._tools[name] = ToolDef(
            name=name,
            description=description,
            input_schema=input_schema,
            handler=handler,
            async_handler=async_handler,
            integration=integration,
            contract_action=contract_action,
            category=category,
            namespace=resolved_ns or "misc",
        )
        if integration and integration in self._integrations:
            self._integrations[integration].tools.append(name)
        self._tool_defs_cache = None  # invalidate
        log.debug(
            "Tool registered: %s (integration: %s, namespace: %s)",
            name, integration or "core", resolved_ns or "misc",
        )

    def register_integration(
        self,
        name: str,
        display_name: str,
        description: str,
        status_fn: Callable[[], bool],
    ) -> None:
        """Register an integration with a live status check."""
        self._integrations[name] = IntegrationDef(
            name=name,
            display_name=display_name,
            description=description,
            status_fn=status_fn,
        )
        self._tool_defs_cache = None  # invalidate
        self._connected_cache = None
        self._capabilities_cache = None
        log.debug("Integration registered: %s", name)

    # ── For LLM ──

    def get_tool_definitions(
        self,
        *,
        namespaces: frozenset[str] | set[str] | None = None,
    ) -> list[dict]:
        """Get tool definitions in the format expected by the LLM.

        Only includes tools whose integration is connected (or core tools).
        Cached for up to 60s to avoid calling status_fn() on every message.

        Args:
            namespaces: when set, only include tools whose namespace is in
                this set (plus external MCP tools and always-misc tools).
                When None, return all tools (existing behavior — default
                everywhere except the orchestrator's namespace-narrowing
                path, which passes an explicit set).
        """
        now = time.monotonic()
        if namespaces is None and self._tool_defs_cache is not None and (now - self._tool_defs_cache_time) < _CACHE_TTL:
            return self._tool_defs_cache

        # Misc is always included so unlisted tools never get silently
        # hidden during rollout. Explicit opt-out would be a future
        # configuration — for now, misc is a safety net.
        ns_filter: set[str] | None = None
        if namespaces is not None:
            ns_filter = set(namespaces) | {"misc"}

        definitions = []
        for tool in self._tools.values():
            # Skip tools whose integration is disconnected
            if tool.integration:
                integration = self._integrations.get(tool.integration)
                if integration and not integration.status_fn():
                    continue
            if ns_filter is not None and tool.namespace not in ns_filter:
                continue
            definitions.append({
                "name": tool.name,
                "description": tool.description,
                "input_schema": tool.input_schema,
            })
        # Append external MCP tools — always included (they have no namespace
        # in our taxonomy and are already opt-in via MCP server config).
        if self._mcp_client:
            for ext in self._mcp_client.get_external_tool_definitions():
                definitions.append({
                    "name": ext["name"],
                    "description": ext.get("description", ""),
                    "input_schema": ext.get("input_schema", {"type": "object", "properties": {}}),
                })
        if namespaces is None:
            self._tool_defs_cache = definitions
            self._tool_defs_cache_time = now
        return definitions

    def get_contract_filtered_definitions(self) -> list[dict]:
        """Get tool definitions excluding contract-blocked tools.

        Unlike ``get_tool_definitions`` (which is integration-aware but
        contract-blind), this method additionally removes any tool whose
        ``contract_action`` is set to *block* in the current CONTRACT.md.
        Use this for the main chat loop so the LLM never wastes a round
        calling a tool the contract will refuse.

        Confirmation-level tools are kept — the user may approve them.
        """
        all_defs = self.get_tool_definitions()
        if not self._contract:
            return all_defs

        filtered: list[dict] = []
        for defn in all_defs:
            tool = self._tools.get(defn["name"])
            if tool and tool.contract_action:
                try:
                    result = self._contract.check(tool.contract_action)
                    if not result.allowed and not result.requires_confirmation:
                        log.debug(
                            "Hiding tool %s from LLM — contract blocks %s",
                            tool.name, tool.contract_action,
                        )
                        continue
                except Exception:
                    pass  # Degrade gracefully — include the tool
            filtered.append(defn)
        return filtered

    def get_capabilities_text(self) -> str:
        """Build the capabilities section for the system prompt (cached 60s)."""
        now = time.monotonic()
        if self._capabilities_cache is not None and (now - self._capabilities_cache_time) < _CACHE_TTL:
            return self._capabilities_cache

        lines = [
            "# Capabilities",
            "",
            "You have tools available. Use them proactively when the user's request "
            "requires action — don't just describe what you could do, actually do it.",
            "",
        ]

        # Group tools by category/integration
        categories: dict[str, list[ToolDef]] = {}
        for tool in self._tools.values():
            key = tool.category or tool.integration or "core"
            categories.setdefault(key, []).append(tool)

        for cat, tools in categories.items():
            # Get integration info if this is an integration category
            integration = self._integrations.get(cat)
            if integration:
                header = f"## {integration.display_name}"
                desc = integration.description
            else:
                header = f"## {cat.title()}"
                desc = ""

            tool_names = ", ".join(t.name for t in tools)
            lines.append(f"{header} (Tools: {tool_names})")
            if desc:
                lines.append(desc)
            for tool in tools:
                lines.append(f"- **{tool.name}**: {tool.description}")
            lines.append("")

        lines.extend([
            "## Reporting",
            "",
            "After completing any action (tool use, task, search, email operation), briefly report:",
            "1. **What you did** — one sentence summary of the action taken",
            "2. **Result** — what happened (success, data found, error)",
            "3. **Next steps** — if there are follow-up actions the user should take or approve",
            "",
            "Keep reports concise. Don't just say \"Done\" — tell the user what was actually done.",
            "",
        ])

        result = "\n".join(lines)
        self._capabilities_cache = result
        self._capabilities_cache_time = now
        return result

    def get_connected_integrations(self) -> list[str]:
        """Get list of currently connected integrations (cached 60s)."""
        now = time.monotonic()
        if self._connected_cache is not None and (now - self._connected_cache_time) < _CACHE_TTL:
            return self._connected_cache

        connected = []
        for integration in self._integrations.values():
            try:
                if integration.status_fn():
                    tool_names = ", ".join(integration.tools) if integration.tools else "no tools"
                    connected.append(f"{integration.display_name} ({tool_names})")
            except Exception:
                pass
        self._connected_cache = connected
        self._connected_cache_time = now
        return connected

    # ── Execution ──

    def execute(self, tool_name: str, tool_input: dict) -> dict[str, Any]:
        """Execute a tool with contract enforcement."""
        tool = self._tools.get(tool_name)
        if not tool:
            return {"error": f"Unknown tool: {tool_name}"}

        # Contract check
        if tool.contract_action and self._contract:
            result = self._contract.check(tool.contract_action)
            if not result.allowed:
                return {"error": f"Contract blocked: {result.reason}"}
            if result.requires_confirmation:
                return {
                    "needs_confirmation": True,
                    "action": tool_name,
                    "reason": result.reason,
                }

        # Execute
        try:
            return tool.handler(tool_input)
        except Exception as e:
            log.exception("Tool execution failed: %s", tool_name)
            return {"error": str(e)}

    # ── Introspection ──

    def list_tools(self) -> list[dict]:
        """List all registered tools (for API/debug)."""
        return [
            {
                "name": t.name,
                "description": t.description,
                "integration": t.integration,
                "category": t.category,
                "contract_action": t.contract_action,
            }
            for t in self._tools.values()
        ]

    def list_integrations(self) -> list[dict]:
        """List all registered integrations with status."""
        result = []
        for i in self._integrations.values():
            try:
                connected = i.status_fn()
            except Exception:
                connected = False
            result.append({
                "name": i.name,
                "display_name": i.display_name,
                "description": i.description,
                "connected": connected,
                "tools": i.tools,
            })
        return result
