"""Tests for the tool registry and AST-based lazy discovery.

Two-layer convention: we exercise the registry directly (no MCP server needed)
and use temp on-disk tool modules where we need to verify the AST scan picks
up real decorators and ignores look-alikes inside strings/comments.
"""

from __future__ import annotations

import ast
import importlib
import sys
import textwrap
from unittest.mock import patch

import pytest

from a360_mcp.utils.registry import (
    _build_tool_module_index,
    _extract_register_tool_name,
    _load_static_tool_module_index,
    clear_registry,
    discover_tools_lazy,
    get_registry,
    get_tool_module_index,
    register_tool,
    reset_discovery_state,
)


@pytest.fixture(autouse=True)
def _reset_registry():
    """Each test starts with an empty registry and unbuilt index."""
    clear_registry()
    reset_discovery_state()
    yield
    clear_registry()
    reset_discovery_state()


# --------------------------------------------------------------------------- #
# register_tool basics
# --------------------------------------------------------------------------- #


class _DummyParams:  # minimal stand-in for a Pydantic model
    pass


def test_register_tool_records_entry():
    @register_tool("foo", params=_DummyParams, description="Foo tool")
    def foo(config, auth, params):
        return {"ok": True}

    assert "foo" in get_registry()
    func, params, ret, desc, ser = get_registry()["foo"]
    assert func is foo
    assert params is _DummyParams
    assert desc == "Foo tool"
    assert ser == "raw_dict"
    assert ret is dict


def test_register_tool_requires_name_and_description():
    with pytest.raises(ValueError):

        @register_tool("", params=_DummyParams, description="d")
        def t1(config, auth, params):
            pass

    with pytest.raises(ValueError):

        @register_tool("ok", params=_DummyParams, description="")
        def t2(config, auth, params):
            pass


def test_register_tool_overwrite_warns():
    @register_tool("dup", params=_DummyParams, description="first")
    def first(config, auth, params):
        pass

    with pytest.warns(UserWarning, match="overwriting"):

        @register_tool("dup", params=_DummyParams, description="second")
        def second(config, auth, params):
            pass

    assert get_registry()["dup"][3] == "second"


def test_register_tool_custom_serialization():
    @register_tool(
        "s", params=_DummyParams, description="d", serialization="columnar", return_type=list
    )
    def s(config, auth, params):
        return []

    _, _, ret, _, ser = get_registry()["s"]
    assert ser == "columnar"
    assert ret is list


# --------------------------------------------------------------------------- #
# AST extraction
# --------------------------------------------------------------------------- #


def _decorator_call(src: str) -> ast.AST:
    tree = ast.parse(textwrap.dedent(src))
    func_def = next(
        n for n in ast.walk(tree) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
    )
    return func_def.decorator_list[0]


def test_extract_name_positional():
    dec = _decorator_call("""
        @register_tool("foo", params=P, description="d")
        def foo(c, a, p): pass
        """)
    assert _extract_register_tool_name(dec) == "foo"


def test_extract_name_keyword():
    dec = _decorator_call("""
        @register_tool(name="bar", params=P, description="d")
        def bar(c, a, p): pass
        """)
    assert _extract_register_tool_name(dec) == "bar"


def test_extract_name_attribute_decorator():
    """``@registry.register_tool(...)`` should also resolve."""
    dec = _decorator_call("""
        @registry.register_tool("baz", params=P, description="d")
        def baz(c, a, p): pass
        """)
    assert _extract_register_tool_name(dec) == "baz"


def test_extract_name_ignores_non_register_decorator():
    dec = _decorator_call("""
        @some_other_decorator("foo")
        def foo(c, a, p): pass
        """)
    assert _extract_register_tool_name(dec) is None


def test_extract_name_ignores_dynamic_name():
    dec = _decorator_call("""
        @register_tool(NAME, params=P, description="d")
        def foo(c, a, p): pass
        """)
    assert _extract_register_tool_name(dec) is None


# --------------------------------------------------------------------------- #
# AST index build (real on-disk package)
# --------------------------------------------------------------------------- #


@pytest.fixture
def fake_tools_pkg(tmp_path, monkeypatch):
    """Create a throwaway tools package on disk with real decorator usage and ghosts."""
    pkg_name = "_fake_a360_tools_for_test"
    pkg_dir = tmp_path / pkg_name
    pkg_dir.mkdir()
    (pkg_dir / "__init__.py").write_text("")
    # Real tool — should be discovered.
    (pkg_dir / "real_tools.py").write_text(textwrap.dedent("""
            from a360_mcp.utils.registry import register_tool

            class P: pass

            @register_tool("real_one", params=P, description="real")
            def real_one(c, a, p):
                return {}

            # Ghost: looks like a decorator but is inside a comment.
            # @register_tool("ghost_in_comment", params=P, description="x")

            GHOST_STRING = "@register_tool('ghost_in_string', params=P, description='x')"

            @register_tool("real_two", params=P, description="real2")
            async def real_two(c, a, p):
                return {}
            """))
    # A module that does NOT use register_tool — must not contribute.
    (pkg_dir / "helpers.py").write_text("def helper():\n    return 1\n")
    # A module starting with _ — must be skipped.
    (pkg_dir / "_internal.py").write_text(
        "from a360_mcp.utils.registry import register_tool\n"
        "class P: pass\n"
        "@register_tool('internal_skipped', params=P, description='x')\n"
        "def f(c,a,p): return {}\n"
    )

    monkeypatch.syspath_prepend(str(tmp_path))
    # Force fresh import.
    for mod in list(sys.modules):
        if mod == pkg_name or mod.startswith(pkg_name + "."):
            del sys.modules[mod]
    importlib.import_module(pkg_name)
    yield pkg_name
    for mod in list(sys.modules):
        if mod == pkg_name or mod.startswith(pkg_name + "."):
            del sys.modules[mod]


def test_build_index_finds_real_tools(fake_tools_pkg):
    index = _build_tool_module_index(fake_tools_pkg)
    assert index["real_one"] == f"{fake_tools_pkg}.real_tools"
    assert index["real_two"] == f"{fake_tools_pkg}.real_tools"


def test_build_index_ignores_ghosts(fake_tools_pkg):
    index = _build_tool_module_index(fake_tools_pkg)
    assert "ghost_in_comment" not in index
    assert "ghost_in_string" not in index
    assert "internal_skipped" not in index


def test_build_index_warns_on_cross_module_dup(fake_tools_pkg, tmp_path):
    pkg_dir = tmp_path / fake_tools_pkg
    (pkg_dir / "other.py").write_text(textwrap.dedent("""
            from a360_mcp.utils.registry import register_tool
            class P: pass
            @register_tool("real_one", params=P, description="dup")
            def real_one(c, a, p): return {}
            """))
    for mod in list(sys.modules):
        if mod == fake_tools_pkg or mod.startswith(fake_tools_pkg + "."):
            del sys.modules[mod]
    importlib.import_module(fake_tools_pkg)
    with pytest.warns(UserWarning, match="defined in both"):
        index = _build_tool_module_index(fake_tools_pkg)
    # First-seen wins; sorted iteration means alphabetical order, so "other"
    # precedes "real_tools".
    assert index["real_one"] == f"{fake_tools_pkg}.other"


# --------------------------------------------------------------------------- #
# Static index + get_tool_module_index caching
# --------------------------------------------------------------------------- #


def test_load_static_index_returns_empty_when_absent(fake_tools_pkg):
    assert _load_static_tool_module_index(fake_tools_pkg) == {}


def test_load_static_index_reads_module(fake_tools_pkg, tmp_path):
    pkg_dir = tmp_path / fake_tools_pkg
    (pkg_dir / "_module_index.py").write_text(
        f'TOOL_MODULE_INDEX = {{"x": "{fake_tools_pkg}.real_tools"}}\n'
    )
    for mod in list(sys.modules):
        if mod == fake_tools_pkg or mod.startswith(fake_tools_pkg + "."):
            del sys.modules[mod]
    importlib.import_module(fake_tools_pkg)
    idx = _load_static_tool_module_index(fake_tools_pkg)
    assert idx == {"x": f"{fake_tools_pkg}.real_tools"}


def test_static_index_wins_over_ast_scan(fake_tools_pkg, tmp_path):
    pkg_dir = tmp_path / fake_tools_pkg
    (pkg_dir / "_module_index.py").write_text(
        f'TOOL_MODULE_INDEX = {{"real_one": "{fake_tools_pkg}.from_static"}}\n'
    )
    for mod in list(sys.modules):
        if mod == fake_tools_pkg or mod.startswith(fake_tools_pkg + "."):
            del sys.modules[mod]
    importlib.import_module(fake_tools_pkg)
    reset_discovery_state()
    index = get_tool_module_index(fake_tools_pkg)
    assert index["real_one"] == f"{fake_tools_pkg}.from_static"


def test_get_tool_module_index_caches(fake_tools_pkg):
    reset_discovery_state()
    first = get_tool_module_index(fake_tools_pkg)
    # Patch the builder to ensure it is not called again.
    with patch(
        "a360_mcp.utils.registry._build_tool_module_index",
        side_effect=AssertionError("should be cached"),
    ):
        second = get_tool_module_index(fake_tools_pkg)
    assert first is second


def test_get_tool_module_index_rebuild_forces_rescan(fake_tools_pkg):
    reset_discovery_state()
    get_tool_module_index(fake_tools_pkg)
    # rebuild=True should call the builder again, not return the cache.
    with patch(
        "a360_mcp.utils.registry._build_tool_module_index",
        return_value={"rebuilt": "x"},
    ) as mock_build:
        # Drop the static index so rebuild path runs the patched builder.
        with patch(
            "a360_mcp.utils.registry._load_static_tool_module_index",
            return_value={},
        ):
            result = get_tool_module_index(fake_tools_pkg, rebuild=True)
    mock_build.assert_called_once()
    assert result == {"rebuilt": "x"}


# --------------------------------------------------------------------------- #
# discover_tools_lazy
# --------------------------------------------------------------------------- #


def test_discover_lazy_none_returns_full_registry(fake_tools_pkg):
    reset_discovery_state()
    out = discover_tools_lazy(None, tools_package=fake_tools_pkg)
    assert "real_one" in out
    assert "real_two" in out


def test_discover_lazy_empty_returns_empty(fake_tools_pkg):
    reset_discovery_state()
    assert discover_tools_lazy([], tools_package=fake_tools_pkg) == {}
    assert discover_tools_lazy(iter(()), tools_package=fake_tools_pkg) == {}


def test_discover_lazy_selective_imports_only_needed_modules(fake_tools_pkg):
    reset_discovery_state()
    # Track which modules get imported.
    real_import = importlib.import_module
    imported: list[str] = []

    def tracking_import(name, *args, **kwargs):
        if name.startswith(fake_tools_pkg):
            imported.append(name)
        return real_import(name, *args, **kwargs)

    with patch("a360_mcp.utils.registry.importlib.import_module", side_effect=tracking_import):
        out = discover_tools_lazy(["real_one"], tools_package=fake_tools_pkg)

    assert "real_one" in out
    assert "real_two" not in out
    # Only the owning module of real_one should be imported (plus the package
    # itself for the AST scan).
    assert f"{fake_tools_pkg}.real_tools" in imported
    assert f"{fake_tools_pkg}.helpers" not in imported


def test_discover_lazy_unknown_name_skipped(fake_tools_pkg):
    reset_discovery_state()
    out = discover_tools_lazy(["does_not_exist"], tools_package=fake_tools_pkg)
    assert out == {}


def test_discover_lazy_partial_known_partial_unknown(fake_tools_pkg):
    reset_discovery_state()
    out = discover_tools_lazy(["real_one", "ghost"], tools_package=fake_tools_pkg)
    assert "real_one" in out
    assert "ghost" not in out


def test_discover_lazy_falls_back_to_eager_when_index_misses(fake_tools_pkg):
    """If the AST index misses a name, full eager discovery runs once."""
    reset_discovery_state()
    # Simulate a tool whose decorator name was dynamic so AST scan missed it.
    # We patch the index to be empty so the fallback path triggers.
    with patch(
        "a360_mcp.utils.registry.get_tool_module_index",
        return_value={},
    ):
        # discover_tools (eager) will import the module, which registers real_one.
        out = discover_tools_lazy(["real_one"], tools_package=fake_tools_pkg)
    assert "real_one" in out
