from __future__ import annotations

import asyncio
import fnmatch
import importlib.metadata
import inspect
import logging
import os
import re
from typing import Any, Protocol, runtime_checkable

from axm_mcp.concurrency import KeyedLock

__all__ = [
    "_HTTP_MODE",
    "_collect_dispatcher_params",
    "_extract_docstring_params",
    "_git_lock",
    "_is_disabled",
    "_log_external_step",
    "_register_one",
    "_session_lock",
    "discover_tools",
    "register_tools",
]

logger = logging.getLogger(__name__)

_EP_GROUP = "axm.tools"

# Set to True when the server runs in HTTP/SSE mode (shared process).
# When True, tools that receive path="." get a warning because "." resolves
# to the server's CWD, not the conversation's workspace.
_HTTP_MODE: bool = False

# Per-key locks — active only in HTTP mode (checked at call time).
_session_lock = KeyedLock()  # protocol_* tools, keyed by session_id
_git_lock = KeyedLock()  # git_* tools, keyed by repo path


def _is_disabled(name: str, patterns: list[str]) -> bool:
    """Check if a tool name matches any disable pattern.

    Supports exact names (``ast_dead_code``) and glob patterns
    (``bib_*``, ``ticket_*``).
    """
    return any(fnmatch.fnmatch(name, pat) for pat in patterns)


@runtime_checkable
class ToolLike(Protocol):
    """Minimal protocol for AXMTool-compatible objects."""

    @property
    def name(self) -> str:
        """Tool name used for MCP registration."""
        ...

    def execute(self, **kwargs: Any) -> Any:
        """Execute the tool with the given keyword arguments."""
        ...


def discover_tools() -> dict[str, Any]:
    """Discover and instantiate all AXMTool entry points.

    Supports both ``AXMTool`` subclasses (instantiated) and plain
    dispatcher functions (used as-is).

    Tools can be excluded via the ``AXM_DISABLE_TOOLS`` environment
    variable — a comma-separated list of tool names or glob patterns
    (e.g. ``bib_*,ticket_*,ast_dead_code``).

    Returns:
        Dict mapping tool name → tool instance or callable.
    """
    raw = os.environ.get("AXM_DISABLE_TOOLS", "")
    disable_patterns = [p.strip() for p in raw.split(",") if p.strip()]
    if disable_patterns:
        logger.info("AXM_DISABLE_TOOLS: %s", disable_patterns)

    tools: dict[str, Any] = {}

    for ep in importlib.metadata.entry_points(group=_EP_GROUP):
        if disable_patterns and _is_disabled(ep.name, disable_patterns):
            logger.info("Skipping disabled tool: %s", ep.name)
            continue
        try:
            obj = ep.load()
            if isinstance(obj, type):
                tool = obj()  # AXMTool class → instantiate
            else:
                tool = obj  # plain function → use as-is
            tools[ep.name] = tool
            logger.debug("Discovered tool: %s", ep.name)
        except Exception:
            logger.warning(
                "Failed to load tool entry point: %s",
                ep.name,
                exc_info=True,
            )

    return tools


def register_tools(
    mcp: Any,
    tools: dict[str, Any],
    extra_tools: dict[str, str] | None = None,
) -> None:
    """Register discovered tools as MCP tool callables.

    Each tool becomes a callable ``tool_name(**kwargs) -> dict``
    that delegates to ``tool.execute(**kwargs)``.

    Args:
        mcp: FastMCP server instance.
        tools: Dict from discover_tools().
        extra_tools: Optional dict of manually-registered tool names
            to their descriptions (for list_tools inclusion).
    """
    for name, tool in tools.items():
        _register_one(mcp, name, tool)
        logger.info("Registered MCP tool: %s", name)


def _find_actions_dict(
    module: Any | None,
) -> dict[str, Any] | None:
    """Find a ``_*_ACTIONS`` dict on *module*."""
    if module is None:
        return None
    for attr_name in dir(module):
        if attr_name.endswith("_ACTIONS"):
            candidate = getattr(module, attr_name, None)
            if isinstance(candidate, dict):
                return candidate
    return None


def _union_subfn_params(
    actions_dict: dict[str, Any],
) -> dict[str, inspect.Parameter]:
    """Collect all typed params from sub-functions, made optional."""
    seen: dict[str, inspect.Parameter] = {}
    for sub_fn in actions_dict.values():
        try:
            sub_sig = inspect.signature(sub_fn)
        except (ValueError, TypeError):
            continue
        for p in sub_sig.parameters.values():
            if p.name == "self" or p.kind == inspect.Parameter.VAR_KEYWORD:
                continue
            if p.name not in seen:
                default = (
                    p.default if p.default is not inspect.Parameter.empty else None
                )
                seen[p.name] = p.replace(default=default)
    return seen


# Regex: matches "    param_name (optional type): description" or
#        "    param_name: description" in Google-style Args blocks.
_ARG_LINE_RE = re.compile(
    r"^\s{4,}(\w+)"  # param name (indented 4+ spaces)
    r"(?:\s*\(([^)]+)\))?"  # optional (type) in parens
    r"\s*:"  # colon separator
    r"\s*(.*)$",  # description (captured but unused)
)

# Map common docstring type hints to Python annotations.
_TYPE_MAP: dict[str, type | None] = {
    "str": str,
    "int": int,
    "float": float,
    "bool": bool,
    "dict": dict,
    "list": list,
    "Path": str,  # paths come as str over MCP
}


def _extract_docstring_params(
    docstring: str | None,
) -> list[inspect.Parameter]:
    """Parse Google-style ``Args:`` section into ``inspect.Parameter`` objects.

    This is the fallback for non-dispatcher tools whose ``execute(**kwargs)``
    has no typed parameters in the signature but documents them in the
    docstring.

    Args:
        docstring: The docstring to parse.

    Returns:
        List of keyword-only ``inspect.Parameter`` with default ``None``.
        Empty list if no ``Args:`` section found.
    """
    if not docstring:
        return []

    lines = docstring.splitlines()
    in_args = False
    params: list[inspect.Parameter] = []

    for line in lines:
        stripped = line.strip()

        # Detect start of Args: block
        if stripped in ("Args:", "Keyword Args:"):
            in_args = True
            continue

        # End of Args: block (next section or blank line after content)
        if in_args and stripped and not stripped.startswith("**") and ":" in stripped:
            # Check if this is a new top-level section (e.g. "Returns:")
            if re.match(r"^[A-Z]\w*:\s*$", stripped):
                break

        if not in_args:
            continue

        # Skip **kwargs line
        if stripped.startswith("**"):
            continue

        # Match param lines
        match = _ARG_LINE_RE.match(line)
        if match:
            name = match.group(1)
            type_hint = match.group(2)

            annotation: type | Any = inspect.Parameter.empty
            if type_hint:
                # Clean up "str, optional" → "str"
                clean_type = type_hint.split(",")[0].strip()
                annotation = _TYPE_MAP.get(clean_type, inspect.Parameter.empty)

            params.append(
                inspect.Parameter(
                    name,
                    kind=inspect.Parameter.KEYWORD_ONLY,
                    default=None,
                    annotation=annotation,
                )
            )

    return params


def _collect_dispatcher_params(
    fn: Any,
    *,
    override_module: Any | None = None,
) -> list[inspect.Parameter] | None:
    """Collect union of typed params from dispatcher sub-functions.

    A *dispatcher* is a function with ``action: str`` + ``**kwargs``
    that routes to sub-functions stored in a module-level ``_*_ACTIONS``
    dict.  This helper introspects all sub-functions and returns the
    union of their parameters (all made optional).

    Args:
        fn: The dispatcher function to introspect.
        override_module: Module to search for ``_*_ACTIONS`` dict.
            If *None*, uses ``inspect.getmodule(fn)``.

    Returns:
        List of ``inspect.Parameter`` if *fn* is a dispatcher, else *None*.
    """
    try:
        sig = inspect.signature(fn)
    except (ValueError, TypeError):
        return None

    params = list(sig.parameters.values())

    # Detect dispatcher pattern: has 'action' param + VAR_KEYWORD
    has_action = any(p.name == "action" for p in params)
    has_varkw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params)
    if not (has_action and has_varkw):
        return None

    # Find the _ACTIONS dict by convention
    module = override_module or inspect.getmodule(fn)
    actions_dict = _find_actions_dict(module)
    if not actions_dict:
        return None

    # Collect + build final: action (required) + sub-fn params (optional)
    seen = _union_subfn_params(actions_dict)
    action_param = next(p for p in params if p.name == "action")
    return [action_param, *sorted(seen.values(), key=lambda p: p.name)]


def _log_external_step(
    tool_name: str,
    tool_args: dict[str, Any],
    success: bool,
    result_str: str,
    duration_ms: int,
) -> None:
    """Log a non-protocol tool call in the active session trace.

    No-op if no session is active, tracing is disabled, or axm-engine
    is not installed. Errors are silently swallowed — tracing must
    never break tool execution.
    """
    try:
        from axm_engine.runtime.orchestrator import get_orchestrator

        orch = get_orchestrator()
        orch.log_external_step(
            tool_name=tool_name,
            tool_args=tool_args,
            result_success=success,
            result_length=len(result_str),
            result_hash="",  # let orchestrator compute
            duration_ms=duration_ms,
        )
    except Exception:  # noqa: S110
        pass  # tracing must never break tool execution


def _register_one(
    mcp: Any,
    name: str,
    tool: Any,
    *,
    override_module: Any | None = None,
) -> None:
    """Register a single tool, capturing in closure.

    Supports both ``AXMTool`` instances (with ``.execute()``) and plain
    dispatcher functions.  Sets the typed signature on the wrapper
    **before** handing it to ``mcp.tool()``, so FastMCP generates the
    correct JSON-Schema for the tool's parameters.

    For *dispatcher* functions (``action`` + ``**kwargs``), introspects
    sub-functions to build a union of all their typed parameters.

    Args:
        mcp: FastMCP server instance.
        name: Tool name for MCP registration.
        tool: Tool instance or plain function.
        override_module: For testing — module to search for ``_*_ACTIONS``.
    """
    import time

    is_plain = callable(tool) and not hasattr(tool, "execute")
    exec_fn: Any = tool if is_plain else tool.execute

    # Protocol tools already trace via orchestrator.run_tool()
    _should_trace = not name.startswith("protocol_")

    def _warn_implicit_path(tool_name: str, kwargs: dict[str, Any]) -> None:
        """Warn when path is '.' or '' in HTTP mode."""
        if _HTTP_MODE and kwargs.get("path") in (".", ""):
            logger.warning(
                "Tool '%s' called with implicit path='.' in HTTP mode. "
                "Pass an explicit absolute path to avoid operating on the "
                "wrong directory.",
                tool_name,
            )

    if is_plain:

        def _wrapper(**kwargs: Any) -> dict[str, Any]:
            # MCP may nest action args inside a "kwargs" key — unwrap.
            if "kwargs" in kwargs and isinstance(kwargs["kwargs"], dict):
                nested = kwargs.pop("kwargs")
                kwargs.update(nested)
            _warn_implicit_path(name, kwargs)
            start_ns = time.perf_counter_ns()
            result: dict[str, Any] = tool(**kwargs)
            duration_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
            if _should_trace:
                try:
                    _log_external_step(name, kwargs, True, str(result), duration_ms)
                except Exception:  # noqa: S110
                    pass
            return result

    else:

        def _wrapper(**kwargs: Any) -> dict[str, Any]:
            # MCP may nest action args inside a "kwargs" key — unwrap.
            if "kwargs" in kwargs and isinstance(kwargs["kwargs"], dict):
                nested = kwargs.pop("kwargs")
                kwargs.update(nested)
            _warn_implicit_path(name, kwargs)
            start_ns = time.perf_counter_ns()
            result = tool.execute(**kwargs)
            duration_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
            output: dict[str, Any] = {"success": result.success, **result.data}
            if result.error:
                output["error"] = result.error
            hint = getattr(result, "hint", None)
            if hint:
                output["hint"] = hint
            if _should_trace:
                try:
                    _log_external_step(
                        name, kwargs, result.success, str(output), duration_ms
                    )
                except Exception:  # noqa: S110
                    pass
            return output

    # Copy docstring.
    _wrapper.__doc__ = exec_fn.__doc__ or f"Execute {name} tool."

    # Wrap with concurrency lock for protocol/git tools in HTTP mode.
    _lock: KeyedLock | None = None
    _key_param: str | None = None
    if name.startswith("protocol_"):
        _lock = _session_lock
        _key_param = "session_id"
    elif name.startswith("git_"):
        _lock = _git_lock
        _key_param = "path"

    if _lock is not None:
        _sync = _wrapper
        _lk = _lock
        _kp = _key_param

        async def _wrapper(**kwargs: Any) -> dict[str, Any]:  # type: ignore[misc]
            if not _HTTP_MODE:
                return _sync(**kwargs)
            key = kwargs.get(_kp)  # type: ignore[arg-type]
            if key is None:
                return _sync(**kwargs)
            async with _lk(key):
                return await asyncio.to_thread(_sync, **kwargs)

        _wrapper.__doc__ = _sync.__doc__

    # For dispatchers (action + **kwargs), build union of sub-fn params.
    # For regular tools, strip 'self' and **kwargs.
    try:
        union_params = _collect_dispatcher_params(
            exec_fn, override_module=override_module
        )
        if union_params is not None:
            params = union_params
        else:
            try:
                exec_sig = inspect.signature(exec_fn, eval_str=True)
            except Exception:
                exec_sig = inspect.signature(exec_fn)
            params = [
                p
                for p in exec_sig.parameters.values()
                if p.name != "self" and p.kind != inspect.Parameter.VAR_KEYWORD
            ]
            # Fallback: if params is empty (e.g. execute(**kwargs)),
            # parse the docstring Args section for parameter info.
            if not params:
                params = _extract_docstring_params(exec_fn.__doc__)
        _wrapper.__signature__ = inspect.Signature(  # type: ignore[attr-defined]
            parameters=params,
            return_annotation=dict[str, Any],
        )
    except (ValueError, TypeError):
        pass  # Fall back to generic **kwargs if introspection fails

    # Register AFTER setting the signature so FastMCP sees it.
    mcp.tool(name=name)(_wrapper)


def _get_tool_doc(tool: Any) -> str:
    """Extract the first-line docstring from a tool or plain callable."""
    if callable(tool) and not hasattr(tool, "execute"):
        doc = tool.__doc__ or ""
    else:
        doc = getattr(tool.execute, "__doc__", None) or ""
    return doc.strip().split("\n")[0]


def _register_list_tools(
    mcp: Any,
    tools: dict[str, Any],
    extra_tools: dict[str, str],
) -> None:
    """Register the list_tools meta-tool."""

    @mcp.tool(name="list_tools")  # type: ignore[untyped-decorator]
    def _list_tools(**kwargs: Any) -> dict[str, Any]:
        """List all available AXM tools with their names and descriptions."""
        tool_list = []
        for name, tool in sorted(tools.items()):
            tool_list.append({"name": name, "description": _get_tool_doc(tool)})
        for name, desc in sorted(extra_tools.items()):
            tool_list.append({"name": name, "description": desc})
        tool_list.sort(key=lambda t: t["name"])
        return {"tools": tool_list, "count": len(tool_list)}

    logger.info("Registered meta-tool: list_tools")
