"""External tool handler registry for matrx-ai.

When matrx-ai is consumed as a package by an external application (e.g. Matrx Local),
some tools in the database have implementations that live in the host application, not
in this codebase. These tools carry ``source_app != "matrx_ai"`` in the database.

The ``ExternalHandlerRegistry`` allows host applications to register async callables
that matrx-ai will dispatch to when the AI model calls one of those tools. The result
flows back through the same guardrails, streaming, persistence, and model-feedback
pipeline as any native tool -- the model cannot tell the difference.

----

There are three integration patterns, from simplest to most flexible:

Pattern 1 — ``ExternalToolAdapter`` (recommended for most apps)
----------------------------------------------------------------
Subclass ``ExternalToolAdapter``, decorate methods with ``@external_tool``, and call
``register()`` once at startup. The adapter auto-dispatches every decorated method::

    from matrx_ai.tools import ExternalToolAdapter, external_tool, ToolContext

    class MyTools(ExternalToolAdapter):
        source_app = "matrx_local"

        @external_tool("screenshot")
        async def take_screenshot(self, args: dict, ctx: ToolContext) -> dict:
            return {"image": await capture_screen(args.get("region"))}

        @external_tool("open_file")
        async def open_file(self, args: dict, ctx: ToolContext) -> dict:
            return {"content": await read_file(args["path"])}

    MyTools().register()   # call once at app startup


Pattern 2 — ``register_external_tool_handler`` (per-tool, plain functions)
---------------------------------------------------------------------------
Register individual async functions, one per tool::

    from matrx_ai.tools import register_external_tool_handler, ToolContext, ToolResult

    async def handle_screenshot(args: dict, ctx: ToolContext) -> ToolResult:
        image_data = await take_local_screenshot(args.get("region"))
        return ToolResult(success=True, output=image_data,
                          tool_name=ctx.tool_name, call_id=ctx.call_id)

    register_external_tool_handler("screenshot", handle_screenshot)


Pattern 3 — ``register_external_app_handler`` (single dispatcher for all tools)
---------------------------------------------------------------------------------
Register one async function that receives every tool call for a source_app::

    from matrx_ai.tools import register_external_app_handler, ToolContext, ToolResult

    async def matrx_local_dispatcher(args: dict, ctx: ToolContext) -> ToolResult:
        result_data = await local_tool_router(ctx.tool_name, args)
        return ToolResult(success=True, output=result_data,
                          tool_name=ctx.tool_name, call_id=ctx.call_id)

    register_external_app_handler("matrx_local", matrx_local_dispatcher)


Resolution order (per execution)
---------------------------------
1. Exact tool-name match  →  Pattern 2 / adapter per-tool handlers
2. App-level fallback     →  Pattern 3 / adapter app-level dispatcher
3. No handler found       →  structured error returned to the model (not a crash)

----

Handler return values
----------------------
Handlers may return any of the following — all are normalised automatically:

- ``ToolResult``          — passed through as-is (preferred)
- ``dict``                — wrapped as ``ToolResult(success=True, output=...)``
                            (if dict has ``"status": "error"`` it becomes a failure)
- ``str`` / any other     — wrapped as ``ToolResult(success=True, output=value)``

----

Register handlers *before* the first AI request that could trigger the tool.
The registry is a process-level singleton; registrations persist for the lifetime
of the process.
"""

from __future__ import annotations

import functools
import traceback
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar

from matrx_utils import vcprint

from matrx_ai.tools.models import ToolContext, ToolError, ToolResult

ExternalToolHandler = Callable[[dict[str, Any], ToolContext], Awaitable[ToolResult]]

_F = TypeVar("_F", bound=Callable[..., Awaitable[Any]])


class ExternalHandlerRegistry:
    """Process-level singleton that maps tool names / source-app names to async callables.

    Thread-safety: Registration is expected to happen at startup (single-threaded);
    reads during request handling are concurrent but safe because dicts in CPython
    have GIL-protected reads.  For truly concurrent registration, callers should
    synchronise externally.
    """

    _instance: ExternalHandlerRegistry | None = None

    def __init__(self) -> None:
        # tool_name → handler (highest priority)
        self._tool_handlers: dict[str, ExternalToolHandler] = {}
        # source_app → catch-all handler for every tool from that app
        self._app_handlers: dict[str, ExternalToolHandler] = {}

    @classmethod
    def get_instance(cls) -> ExternalHandlerRegistry:
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------

    def register(self, tool_name: str, handler: ExternalToolHandler) -> None:
        """Register a handler for one specific tool by name.

        Takes priority over any app-level handler for the same tool.
        Calling this again with the same ``tool_name`` replaces the previous handler.
        """
        if not callable(handler):
            raise TypeError(f"Handler for tool '{tool_name}' must be callable, got {type(handler)}")
        self._tool_handlers[tool_name] = handler
        vcprint(
            f"[ExternalHandlerRegistry] Registered handler for tool: {tool_name}",
            color="cyan",
        )

    def register_app_handler(self, source_app: str, handler: ExternalToolHandler) -> None:
        """Register a catch-all handler for every tool owned by ``source_app``.

        This is the simplest integration point: one async function receives every
        tool call from the given app, with ``ctx.tool_name`` identifying which
        specific tool was invoked.

        Calling this again with the same ``source_app`` replaces the previous handler.
        """
        if not callable(handler):
            raise TypeError(
                f"App handler for source_app '{source_app}' must be callable, got {type(handler)}"
            )
        self._app_handlers[source_app] = handler
        vcprint(
            f"[ExternalHandlerRegistry] Registered app-level handler for source_app: {source_app}",
            color="cyan",
        )

    def unregister(self, tool_name: str) -> bool:
        """Remove the per-tool handler for ``tool_name``. Returns True if one existed."""
        existed = tool_name in self._tool_handlers
        self._tool_handlers.pop(tool_name, None)
        return existed

    def unregister_app_handler(self, source_app: str) -> bool:
        """Remove the app-level handler for ``source_app``. Returns True if one existed."""
        existed = source_app in self._app_handlers
        self._app_handlers.pop(source_app, None)
        return existed

    # ------------------------------------------------------------------
    # Resolution
    # ------------------------------------------------------------------

    def get_handler(
        self, tool_name: str, source_app: str | None
    ) -> ExternalToolHandler | None:
        """Resolve the best handler for a tool call.

        Resolution order:
        1. Exact tool-name match (``register`` / ``register_external_tool_handler``)
        2. App-level fallback (``register_app_handler`` / ``register_external_app_handler``)
        3. ``None`` — caller should surface a structured error to the model
        """
        if tool_name in self._tool_handlers:
            return self._tool_handlers[tool_name]
        if source_app and source_app in self._app_handlers:
            return self._app_handlers[source_app]
        return None

    def has_handler(self, tool_name: str, source_app: str | None) -> bool:
        """Return True if any handler would be found for this tool."""
        return self.get_handler(tool_name, source_app) is not None

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    @property
    def registered_tools(self) -> list[str]:
        return list(self._tool_handlers.keys())

    @property
    def registered_apps(self) -> list[str]:
        return list(self._app_handlers.keys())

    def __repr__(self) -> str:
        return (
            f"ExternalHandlerRegistry("
            f"tools={self.registered_tools}, "
            f"apps={self.registered_apps})"
        )


# ------------------------------------------------------------------
# Decorator for use with ExternalToolAdapter
# ------------------------------------------------------------------

def external_tool(tool_name: str) -> Callable[[_F], _F]:
    """Decorator that marks an ``ExternalToolAdapter`` method as the handler for a tool.

    Args:
        tool_name: The exact tool name as stored in the ``tools`` database table.

    Usage::

        class MyTools(ExternalToolAdapter):
            source_app = "matrx_local"

            @external_tool("screenshot")
            async def take_screenshot(self, args: dict, ctx: ToolContext) -> dict:
                return {"image": await capture_screen(args.get("region"))}

    The method may return any value — ``ToolResult``, ``dict``, ``str``, or anything
    else.  The adapter normalises the return value before it reaches the model.

    The ``self`` argument receives the adapter instance, giving methods access to any
    shared state or configuration stored on the class.
    """
    def decorator(func: _F) -> _F:
        func._external_tool_name = tool_name  # type: ignore[attr-defined]

        @functools.wraps(func)
        async def wrapper(self: ExternalToolAdapter, args: dict[str, Any], ctx: ToolContext) -> Any:
            return await func(self, args, ctx)

        wrapper._external_tool_name = tool_name  # type: ignore[attr-defined]
        return wrapper  # type: ignore[return-value]

    return decorator


# ------------------------------------------------------------------
# High-level adapter class
# ------------------------------------------------------------------

class ExternalToolAdapter:
    """High-level base class for integrating host-application tools with matrx-ai.

    Subclass this, set ``source_app``, decorate methods with ``@external_tool``,
    then call ``register()`` once at application startup.  The adapter auto-discovers
    all decorated methods and registers them — no manual per-tool registration needed.

    The adapter registers in two ways simultaneously:
    - Each ``@external_tool``-decorated method is registered as a per-tool handler
      (highest priority, exact name match).
    - The adapter itself is registered as the app-level fallback dispatcher
      (catches any tool from ``source_app`` that has no specific handler).

    This means you can mix decorated methods with undecorated ones for custom routing
    by overriding ``dispatch()``.

    Minimal example::

        from matrx_ai.tools import ExternalToolAdapter, external_tool, ToolContext

        class MatrxLocalTools(ExternalToolAdapter):
            source_app = "matrx_local"

            @external_tool("screenshot")
            async def take_screenshot(self, args: dict, ctx: ToolContext) -> dict:
                image = await capture_screen(args.get("region"))
                return {"image_base64": image}

            @external_tool("open_file")
            async def open_file(self, args: dict, ctx: ToolContext) -> dict:
                content = await read_file(args["path"])
                return {"content": content}

        MatrxLocalTools().register()

    Advanced — custom dispatch for unregistered tools::

        class MatrxLocalTools(ExternalToolAdapter):
            source_app = "matrx_local"

            @external_tool("screenshot")
            async def take_screenshot(self, args: dict, ctx: ToolContext) -> dict:
                ...

            async def dispatch(self, args: dict, ctx: ToolContext) -> ToolResult:
                # Called for any tool from source_app NOT covered by @external_tool.
                # ctx.tool_name tells you which tool was invoked.
                result = await self._route_to_local_impl(ctx.tool_name, args)
                return ToolResult(success=True, output=result,
                                  tool_name=ctx.tool_name, call_id=ctx.call_id)
    """

    #: Must be set on the subclass to the value of the ``source_app`` DB column.
    source_app: str = ""

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        # Collect all @external_tool-decorated methods at class definition time.
        cls._tool_methods: dict[str, str] = {}
        for attr_name in dir(cls):
            method = getattr(cls, attr_name, None)
            tool_name = getattr(method, "_external_tool_name", None)
            if tool_name is not None:
                cls._tool_methods[tool_name] = attr_name

    def register(self, registry: ExternalHandlerRegistry | None = None) -> None:
        """Register all decorated methods, the app-level dispatcher, and the
        conversation-cleanup callback with matrx-ai's ``ToolLifecycleManager``.

        Call once at application startup, after ``matrx_ai.initialize()``.

        Args:
            registry: Override the registry instance (defaults to the global singleton).
                      Useful in tests.
        """
        if not self.source_app:
            raise ValueError(
                f"{type(self).__name__}.source_app must be set to the database source_app value "
                "(e.g. 'matrx_local') before calling register()."
            )

        reg = registry or ExternalHandlerRegistry.get_instance()

        # Register each decorated method as a per-tool handler.
        for tool_name, method_name in self.__class__._tool_methods.items():
            method = getattr(self, method_name)
            reg.register(tool_name, self._make_handler(method, tool_name))

        # Register the adapter itself as the app-level fallback.
        reg.register_app_handler(self.source_app, self._app_dispatcher)

        # Wire on_conversation_end into ToolLifecycleManager so conversation-scoped
        # resources (session pools, caches, etc.) are cleaned up automatically.
        try:
            from matrx_ai.tools.lifecycle import ToolLifecycleManager

            lifecycle = ToolLifecycleManager.get_instance()
            lifecycle.register_external_adapter_cleanup(self._on_conversation_end)
        except Exception:
            pass  # lifecycle is optional — adapter still functions without it

        vcprint(
            {
                "source_app": self.source_app,
                "registered_tools": list(self.__class__._tool_methods.keys()),
            },
            f"[ExternalToolAdapter] {type(self).__name__} registered",
            color="cyan",
        )

    def _make_handler(
        self, method: Callable[..., Awaitable[Any]], tool_name: str
    ) -> ExternalToolHandler:
        """Wrap a decorated method as a normalised ``ExternalToolHandler``."""
        async def handler(args: dict[str, Any], ctx: ToolContext) -> ToolResult:
            import time

            started_at = time.time()
            try:
                raw = await method(args, ctx)
            except Exception as exc:
                return ToolResult(
                    success=False,
                    error=ToolError(
                        error_type="external_handler_error",
                        message=str(exc),
                        traceback=traceback.format_exc(),
                        is_retryable=False,
                        suggested_action="The tool implementation raised an unexpected exception.",
                    ),
                    started_at=started_at,
                    completed_at=time.time(),
                    tool_name=tool_name,
                    call_id=ctx.call_id,
                )
            return _normalise_result(raw, tool_name, ctx.call_id, started_at)

        return handler

    async def _app_dispatcher(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult:
        """App-level fallback: routes to ``dispatch()`` for tools not covered by ``@external_tool``."""
        return await self.dispatch(args, ctx)

    async def _on_conversation_end(self, conversation_id: str) -> None:
        """Called by ToolLifecycleManager when a conversation ends or times out.

        Routes to ``on_conversation_end()`` which subclasses should override to
        clean up conversation-scoped resources (session pools, open connections, etc.).
        """
        try:
            await self.on_conversation_end(conversation_id)
        except Exception as exc:
            vcprint(
                f"[ExternalToolAdapter] on_conversation_end raised for conversation "
                f"'{conversation_id}': {exc}",
                color="red",
            )

    async def on_conversation_end(self, conversation_id: str) -> None:
        """Override to clean up conversation-scoped resources when a conversation ends.

        This is called automatically by matrx-ai's ``ToolLifecycleManager`` when a
        conversation ends (either explicitly via ``cleanup_conversation()`` or via the
        idle-timeout sweep after 30 minutes of inactivity).

        Use this to release any state keyed by ``conversation_id`` — session pools,
        open browser instances, background processes, file handles, etc.

        Example::

            class MatrxLocalTools(ExternalToolAdapter):
                source_app = "matrx_local"

                def __init__(self):
                    self._sessions: dict[str, ToolSession] = {}

                async def on_conversation_end(self, conversation_id: str) -> None:
                    session = self._sessions.pop(conversation_id, None)
                    if session:
                        await session.cleanup()
        """

    async def dispatch(self, args: dict[str, Any], ctx: ToolContext) -> ToolResult:
        """Override to handle tools from ``source_app`` that have no ``@external_tool`` handler.

        The default implementation returns an informative error so the model knows the
        tool is not implemented in this adapter.  Override for custom routing logic.
        """
        import time

        decorated = list(self.__class__._tool_methods.keys())
        return ToolResult(
            success=False,
            error=ToolError(
                error_type="not_implemented",
                message=(
                    f"Tool '{ctx.tool_name}' (source_app='{self.source_app}') has no handler "
                    f"in {type(self).__name__}. "
                    f"Decorated tools: {decorated or 'none'}. "
                    f"Override dispatch() or add @external_tool('{ctx.tool_name}') to a method."
                ),
                is_retryable=False,
                suggested_action=(
                    "This tool is not available in the current environment. "
                    "Inform the user and suggest an alternative approach."
                ),
            ),
            started_at=time.time(),
            completed_at=time.time(),
            tool_name=ctx.tool_name,
            call_id=ctx.call_id,
        )


# ------------------------------------------------------------------
# Internal normalisation helper (shared by adapter + invoke_external_handler)
# ------------------------------------------------------------------

def _normalise_result(
    raw: Any,
    tool_name: str,
    call_id: str,
    started_at: float,
) -> ToolResult:
    """Convert any handler return value to a ``ToolResult``."""
    import time

    if isinstance(raw, ToolResult):
        raw.started_at = raw.started_at or started_at
        raw.completed_at = raw.completed_at or time.time()
        raw.tool_name = raw.tool_name or tool_name
        raw.call_id = raw.call_id or call_id
        return raw

    if isinstance(raw, dict) and raw.get("status") == "error":
        return ToolResult(
            success=False,
            error=ToolError(
                error_type="execution",
                message=str(raw.get("error", raw.get("result", "Unknown error"))),
            ),
            started_at=started_at,
            completed_at=time.time(),
            tool_name=tool_name,
            call_id=call_id,
        )

    return ToolResult(
        success=True,
        output=raw,
        started_at=started_at,
        completed_at=time.time(),
        tool_name=tool_name,
        call_id=call_id,
    )


# ------------------------------------------------------------------
# Module-level helpers — the primary public API for host applications
# ------------------------------------------------------------------

def register_external_tool_handler(
    tool_name: str, handler: ExternalToolHandler
) -> None:
    """Register an async handler for a single named tool.

    The handler signature must be::

        async def my_handler(args: dict, ctx: ToolContext) -> ToolResult: ...

    Args:
        tool_name: Exact name of the tool as stored in the ``tools`` database table.
        handler:   Async callable that executes the tool and returns a ``ToolResult``.

    This registration takes priority over any app-level handler registered via
    ``register_external_app_handler``.
    """
    ExternalHandlerRegistry.get_instance().register(tool_name, handler)


def register_external_app_handler(
    source_app: str, handler: ExternalToolHandler
) -> None:
    """Register a catch-all async handler for all tools owned by ``source_app``.

    The handler receives every tool call for tools whose ``source_app`` column in
    the database matches the given string.  Use ``ctx.tool_name`` inside the handler
    to route to the correct local implementation.

    The handler signature must be::

        async def my_dispatcher(args: dict, ctx: ToolContext) -> ToolResult: ...

    Args:
        source_app: Value of the ``source_app`` DB column for the tools this handles
                    (e.g. ``"matrx_local"``).
        handler:    Async callable that dispatches to the appropriate local tool.
    """
    ExternalHandlerRegistry.get_instance().register_app_handler(source_app, handler)


async def invoke_external_handler(
    tool_name: str,
    source_app: str | None,
    args: dict[str, Any],
    ctx: ToolContext,
) -> ToolResult:
    """Invoke the registered handler for a tool, returning a normalised ``ToolResult``.

    This is called internally by ``ToolExecutor._execute_external_handler``.
    It handles:
    - handler lookup (with graceful error if none registered)
    - invocation with timeout guard inherited from the executor
    - normalisation of raw return values into ``ToolResult``
    - exception capture → structured ``ToolResult`` with error details
    """
    import time

    registry = ExternalHandlerRegistry.get_instance()
    handler = registry.get_handler(tool_name, source_app)

    if handler is None:
        app_label = f"'{source_app}'" if source_app else "unknown"
        vcprint(
            {
                "tool_name": tool_name,
                "source_app": source_app,
                "registered_tools": registry.registered_tools,
                "registered_apps": registry.registered_apps,
            },
            f"[ExternalHandlerRegistry] No handler registered for tool '{tool_name}' (source_app={app_label})",
            color="red",
        )
        return ToolResult(
            success=False,
            error=ToolError(
                error_type="no_handler",
                message=(
                    f"Tool '{tool_name}' (source_app={app_label}) requires an external handler "
                    f"that has not been registered. "
                    f"Call register_external_tool_handler('{tool_name}', handler) or "
                    f"register_external_app_handler({app_label}, handler) at startup."
                ),
                is_retryable=False,
                suggested_action=(
                    "This tool cannot be executed in the current environment. "
                    "Inform the user and suggest an alternative approach."
                ),
            ),
            started_at=time.time(),
            completed_at=time.time(),
            tool_name=tool_name,
            call_id=ctx.call_id,
        )

    started_at = time.time()
    try:
        raw_result = await handler(args, ctx)
    except Exception as exc:
        return ToolResult(
            success=False,
            error=ToolError(
                error_type="external_handler_error",
                message=str(exc),
                traceback=traceback.format_exc(),
                is_retryable=False,
                suggested_action="The external tool handler raised an unexpected exception. Check the host application logs.",
            ),
            started_at=started_at,
            completed_at=time.time(),
            tool_name=tool_name,
            call_id=ctx.call_id,
        )

    return _normalise_result(raw_result, tool_name, ctx.call_id, started_at)
