"""Tool registry with AST-based lazy discovery.

A domain tool is a function decorated with ``@register_tool``. The decorator
records the function and its schema in a process-global registry; nothing is
imported eagerly. Two consequences drive this module's design:

1. **Lean packages.** The server ships tools in layers
   (``core`` → ``bot_read``/``log_read`` → ``diagnose`` → ``full``) and only
   imports the modules the active package needs. Importing every tool module
   on startup pulls in heavy transitive imports (HTTP clients, browser
   drivers) for tools the user will never call — that tax compounds at every
   request because the LLM sees the tool list each turn.
2. **No ghost tools.** A naive ``grep`` for ``register_tool`` over the source
   tree matches decorator names inside comments, docstrings and string
   literals. We walk the parsed AST instead, so only real decorator calls on
   real function definitions produce entries.

Discovery has two paths and prefers the cheaper one:

* A pre-generated ``tools/_module_index.py`` (a literal ``{tool_name: module}``
  map) is consulted first. It is the fastest path and survives renames as long
  as someone regenerates it.
* If the static index is missing or empty, we scan the tools package directory
  with ``ast`` (no imports), build the same map, and cache it.

``discover_tools_lazy`` then imports only the modules that contain an enabled
tool name, which is the actual cost saver.
"""

from __future__ import annotations

import ast
import importlib
import logging
import os
import pkgutil
import warnings
from collections.abc import Callable, Iterable
from typing import Any, TypeVar

logger = logging.getLogger(__name__)

# --------------------------------------------------------------------------- #
# Registry types
# --------------------------------------------------------------------------- #

# (tool_func, params_model, return_type, description, serialization)
ToolEntry = tuple[Callable[..., Any], type, type, str, str]

_REGISTRY_LOCK_NOT_NEEDED: bool = True  # registry mutated only at import time

_TOOL_REGISTRY: dict[str, ToolEntry] = {}
_TOOLS_DISCOVERED = False

# tool_name -> module_dotted_path. ``None`` means "not built yet"; ``{}`` means
# "built and empty" (which is distinct from "not yet attempted").
_TOOL_MODULE_INDEX: dict[str, str] | None = None

ParamsT = TypeVar("ParamsT")

DEFAULT_SERIALIZATION = "raw_dict"


def register_tool(
    name: str,
    *,
    params: type,
    description: str,
    serialization: str = DEFAULT_SERIALIZATION,
    return_type: type = dict,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """Decorator: record a tool function under ``name``.

    Re-registering an existing name overwrites the prior entry and emits a
    warning, because two tools with the same name is almost always a typo or a
    rename that forgot to delete the old definition — silent overwrite would
    let the wrong one win depending on import order.
    """
    if not name or not isinstance(name, str):
        raise ValueError("register_tool: name must be a non-empty string")
    if not description or not isinstance(description, str):
        raise ValueError("register_tool: description must be a non-empty string")

    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        if name in _TOOL_REGISTRY:
            warnings.warn(
                f"register_tool: overwriting existing tool {name!r} "
                f"(previous: {_TOOL_REGISTRY[name][0].__module__}.{_TOOL_REGISTRY[name][0].__name__})",
                stacklevel=2,
            )
        _TOOL_REGISTRY[name] = (func, params, return_type, description, serialization)
        # Eagerly record the owning module so lazy discovery can skip a full
        # scan when the decorator has already run for every enabled tool.
        if _TOOL_MODULE_INDEX is not None and name not in _TOOL_MODULE_INDEX:
            _TOOL_MODULE_INDEX[name] = func.__module__
        return func

    return decorator


def get_registry() -> dict[str, ToolEntry]:
    """Return a shallow copy of the current registry (does not discover)."""
    return dict(_TOOL_REGISTRY)


def clear_registry() -> None:
    """Test-only: wipe the registry and discovery state."""
    global _TOOLS_DISCOVERED, _TOOL_MODULE_INDEX
    _TOOL_REGISTRY.clear()
    _TOOLS_DISCOVERED = False
    _TOOL_MODULE_INDEX = None


def reset_discovery_state() -> None:
    """Forget that discovery ran, without dropping registered tools.

    Useful for tests that swap the static index and want the next
    ``discover_tools_lazy`` call to rebuild from scratch.
    """
    global _TOOLS_DISCOVERED, _TOOL_MODULE_INDEX
    _TOOLS_DISCOVERED = False
    _TOOL_MODULE_INDEX = None


# --------------------------------------------------------------------------- #
# Full eager discovery (fallback / cold start)
# --------------------------------------------------------------------------- #


def discover_tools(tools_package: str = "a360_mcp.tools") -> dict[str, ToolEntry]:
    """Import every module in ``tools_package`` so all decorators run.

    Used when no static index is available and the AST scan also fails (e.g.
    tools package shipped as a zip). Best-effort: import errors are warned and
    skipped so one broken tool module does not poison the whole registry.
    """
    global _TOOLS_DISCOVERED
    try:
        pkg = importlib.import_module(tools_package)
    except Exception as exc:  # pragma: no cover - defensive
        logger.warning("discover_tools: cannot import %s: %s", tools_package, exc)
        _TOOLS_DISCOVERED = True
        return dict(_TOOL_REGISTRY)

    pkg_path = getattr(pkg, "__path__", None)
    if pkg_path is None:  # pragma: no cover - defensive
        _TOOLS_DISCOVERED = True
        return dict(_TOOL_REGISTRY)

    for mod_info in pkgutil.iter_modules(pkg_path):
        full_name = f"{tools_package}.{mod_info.name}"
        if mod_info.name.startswith("_"):
            continue
        try:
            importlib.import_module(full_name)
        except Exception as exc:  # pragma: no cover - defensive per-module
            logger.warning("discover_tools: failed to import %s: %s", full_name, exc)

    _TOOLS_DISCOVERED = True
    return dict(_TOOL_REGISTRY)


# --------------------------------------------------------------------------- #
# AST-based static index (the cheap path)
# --------------------------------------------------------------------------- #


def _extract_register_tool_name(dec: ast.AST) -> str | None:
    """Return the ``name`` argument of a ``@register_tool(...)`` call.

    Accepts both positional (``@register_tool("foo")``) and keyword
    (``@register_tool(name="foo")``) forms. Returns ``None`` for anything that
    is not a static string constant — dynamic names are out of scope for a
    static index and the eager path will pick them up instead.
    """
    if not isinstance(dec, ast.Call):
        return None
    func = dec.func
    is_register = False
    if isinstance(func, ast.Name) and func.id == "register_tool":
        is_register = True
    elif isinstance(func, ast.Attribute) and func.attr == "register_tool":
        is_register = True
    if not is_register:
        return None

    if dec.args:
        first = dec.args[0]
        if isinstance(first, ast.Constant) and isinstance(first.value, str):
            return first.value
    for kw in dec.keywords:
        if kw.arg == "name" and isinstance(kw.value, ast.Constant):
            if isinstance(kw.value.value, str):
                return kw.value.value
    return None


def _build_tool_module_index(
    tools_package: str = "a360_mcp.tools",
) -> dict[str, str]:
    """Walk the tools package with ``ast`` and map tool_name -> module path.

    No imports happen here. The scan is intentionally cheap so the server can
    decide which modules to load on a cold start without paying for the import
    of every tool module.
    """
    try:
        pkg = importlib.import_module(tools_package)
    except Exception as exc:  # pragma: no cover - defensive
        logger.warning("_build_tool_module_index: cannot import %s: %s", tools_package, exc)
        return {}

    pkg_path = getattr(pkg, "__path__", None)
    if not pkg_path:  # pragma: no cover - defensive
        return {}

    # Resolve the on-disk directory; iter_modules would also work but ast needs
    # file contents, so we read the directory directly.
    pkg_dir = (
        os.path.dirname(next(iter(pkg_path)).__file__ or "")
        if hasattr(next(iter(pkg_path)), "__file__")
        else None
    )
    # The above is awkward; fall back to pkg.__file__ for the common case.
    pkg_file = getattr(pkg, "__file__", None)
    if pkg_file:
        pkg_dir = os.path.dirname(pkg_file)
    if not pkg_dir or not os.path.isdir(pkg_dir):
        return {}

    index: dict[str, str] = {}
    for entry in sorted(os.listdir(pkg_dir)):
        if not entry.endswith(".py"):
            continue
        if entry.startswith("_"):
            continue
        mod_name = entry[:-3]
        full_name = f"{tools_package}.{mod_name}"
        file_path = os.path.join(pkg_dir, entry)
        try:
            with open(file_path, "rb") as fh:
                tree = ast.parse(fh.read(), filename=file_path)
        except (OSError, SyntaxError) as exc:  # pragma: no cover - defensive
            logger.warning("_build_tool_module_index: cannot parse %s: %s", file_path, exc)
            continue

        for node in ast.walk(tree):
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            for dec in node.decorator_list:
                tool_name = _extract_register_tool_name(dec)
                if tool_name is None:
                    continue
                if tool_name in index and index[tool_name] != full_name:
                    warnings.warn(
                        f"_build_tool_module_index: tool {tool_name!r} defined in both "
                        f"{index[tool_name]} and {full_name}; keeping the first.",
                        stacklevel=2,
                    )
                    continue
                index.setdefault(tool_name, full_name)
    return index


def _load_static_tool_module_index(
    tools_package: str = "a360_mcp.tools",
) -> dict[str, str]:
    """Import ``tools/_module_index.py`` if present and return its map."""
    mod_name = f"{tools_package}._module_index"
    try:
        mod = importlib.import_module(mod_name)
    except ModuleNotFoundError:
        return {}
    except Exception as exc:  # pragma: no cover - defensive
        logger.warning("_load_static_tool_module_index: %s: %s", mod_name, exc)
        return {}
    raw = getattr(mod, "TOOL_MODULE_INDEX", None)
    if not isinstance(raw, dict):
        return {}
    # Keep only string->string entries; anything else is a corrupt index.
    return {str(k): str(v) for k, v in raw.items() if isinstance(k, str) and isinstance(v, str)}


def get_tool_module_index(
    tools_package: str = "a360_mcp.tools",
    *,
    rebuild: bool = False,
) -> dict[str, str]:
    """Return the cached tool_name -> module map, building it on first use.

    Static index wins when present and non-empty; otherwise the AST scan
    result is cached. ``rebuild=True`` forces a fresh AST scan (the static
    index, if any, is still tried first because it represents an explicit
    authoring choice).
    """
    global _TOOL_MODULE_INDEX
    if _TOOL_MODULE_INDEX is not None and not rebuild:
        return _TOOL_MODULE_INDEX

    static = _load_static_tool_module_index(tools_package)
    if static:
        _TOOL_MODULE_INDEX = static
        return static

    _TOOL_MODULE_INDEX = _build_tool_module_index(tools_package)
    return _TOOL_MODULE_INDEX


# --------------------------------------------------------------------------- #
# Lazy discovery — the path the server actually uses
# --------------------------------------------------------------------------- #


def discover_tools_lazy(
    enabled_names: Iterable[str] | None,
    *,
    tools_package: str = "a360_mcp.tools",
) -> dict[str, ToolEntry]:
    """Return registry entries for ``enabled_names``, importing only what is needed.

    * ``enabled_names is None`` → discover everything (full eager import).
    * Empty iterable → empty dict (the server chose to expose no tools).
    * Otherwise → resolve the modules that own the enabled names via the
      static/AST index and import only those.

    Unknown names are silently skipped here; the server is responsible for
    telling the LLM which tools it advertised, so a name not in the index is
    either a typo or a future tool, neither of which should crash discovery.
    """
    global _TOOLS_DISCOVERED

    if enabled_names is None:
        if not _TOOLS_DISCOVERED:
            discover_tools(tools_package)
        return dict(_TOOL_REGISTRY)

    wanted = [n for n in enabled_names if n]
    if not wanted:
        return {}

    index = get_tool_module_index(tools_package)
    modules_needed: set[str] = set()
    missing: list[str] = []
    for name in wanted:
        mod = index.get(name)
        if mod is None:
            missing.append(name)
            continue
        modules_needed.add(mod)

    if missing:
        # Fall back to a full eager discovery once. This catches tools whose
        # decorator used a dynamic name, or a stale static index. After this,
        # the registry is authoritative.
        if not _TOOLS_DISCOVERED:
            discover_tools(tools_package)
        # Anything still missing genuinely does not exist.
        still_missing = [n for n in missing if n not in _TOOL_REGISTRY]
        if still_missing:
            logger.debug(
                "discover_tools_lazy: enabled names not found in any module: %s",
                still_missing,
            )

    for mod_dotted in sorted(modules_needed):
        try:
            importlib.import_module(mod_dotted)
        except Exception as exc:  # pragma: no cover - defensive per-module
            logger.warning("discover_tools_lazy: cannot import %s: %s", mod_dotted, exc)

    _TOOLS_DISCOVERED = True
    return {name: _TOOL_REGISTRY[name] for name in wanted if name in _TOOL_REGISTRY}
