# SPDX-License-Identifier: Apache-2.0
# Standard
from types import ModuleType
from unittest.mock import AsyncMock, patch

# Third Party
import pytest

# First Party
from lmcache.v1.basic_check import main, parse_args
from lmcache.v1.check import CheckModeRegistry, check_mode

# --------------- CheckModeRegistry tests ---------------


class TestCheckModeRegistry:
    """Tests for CheckModeRegistry."""

    def test_register_and_get(self) -> None:
        reg = CheckModeRegistry()
        reg.loaded = True  # skip auto-loading

        async def dummy(**kw):
            pass

        reg.register("foo", dummy)
        assert reg.get_mode("foo") is dummy

    def test_register_duplicate_raises(self) -> None:
        reg = CheckModeRegistry()
        reg.loaded = True

        reg.register("dup", lambda: None)
        with pytest.raises(ValueError, match="already registered"):
            reg.register("dup", lambda: None)

    def test_get_mode_returns_none_for_unknown(self) -> None:
        reg = CheckModeRegistry()
        reg.loaded = True
        assert reg.get_mode("nonexistent") is None

    def test_load_modes_discovers_decorated_functions(
        self, tmp_path, monkeypatch
    ) -> None:
        """Verify that load_modes discovers check_mode_ modules."""
        reg = CheckModeRegistry()

        # Create a fake check_mode_ module
        mod = ModuleType("check_mode_fake")

        async def fake_func(**kw):
            pass

        fake_func.is_check_mode = True  # type: ignore[attr-defined]
        fake_func.mode_name = "fake"  # type: ignore[attr-defined]
        mod.fake_func = fake_func  # type: ignore[attr-defined]

        # Patch os.listdir to return our fake module
        monkeypatch.setattr(
            "os.listdir",
            lambda _: ["check_mode_fake.py", "utils.py"],
        )
        monkeypatch.setattr(
            "importlib.import_module",
            lambda name, package=None: mod,
        )

        reg.load_modes()
        assert reg.loaded is True
        assert "fake" in reg.modes
        assert reg.modes["fake"] is fake_func

    def test_load_modes_skips_import_errors(self, monkeypatch) -> None:
        """load_modes should not crash on ImportError."""
        reg = CheckModeRegistry()

        def raiser(name, package=None):
            raise ImportError("boom")

        monkeypatch.setattr(
            "os.listdir",
            lambda _: ["check_mode_bad.py"],
        )
        monkeypatch.setattr("importlib.import_module", raiser)

        reg.load_modes()
        assert reg.loaded is True
        assert len(reg.modes) == 0

    def test_load_modes_called_once(self, monkeypatch) -> None:
        """Second call to load_modes should be a no-op."""
        reg = CheckModeRegistry()

        call_count = 0

        def counting_listdir(path):
            nonlocal call_count
            call_count += 1
            return []

        monkeypatch.setattr("os.listdir", counting_listdir)

        reg.load_modes()
        reg.load_modes()
        assert call_count == 1

    def test_get_mode_triggers_load(self, monkeypatch) -> None:
        """get_mode should call load_modes if not yet loaded."""
        reg = CheckModeRegistry()
        assert reg.loaded is False

        monkeypatch.setattr("os.listdir", lambda _: [])

        reg.get_mode("any")
        assert reg.loaded is True


# --------------- check_mode decorator tests ---------------


class TestCheckModeDecorator:
    """Tests for the @check_mode decorator."""

    def test_decorator_sets_attributes(self) -> None:
        @check_mode("my_mode")
        async def handler(**kw):
            pass

        assert handler.is_check_mode is True  # type: ignore[attr-defined]
        assert handler.mode_name == "my_mode"  # type: ignore[attr-defined]

    def test_decorated_function_still_callable(self) -> None:
        @check_mode("callable_check")
        def sync_fn():
            return 42

        assert sync_fn() == 42


# --------------- parse_args tests ---------------


class TestParseArgs:
    """Tests for basic_check.parse_args."""

    def test_required_mode(self) -> None:
        with patch(
            "sys.argv",
            ["basic_check", "--mode", "test_remote"],
        ):
            args = parse_args()
        assert args.mode == "test_remote"

    @pytest.mark.parametrize(
        "flag,value,attr,expected",
        [
            ("--model", "/my/model", "model", "/my/model"),
            ("--num-keys", "50", "num_keys", 50),
            ("--concurrency", "8", "concurrency", 8),
            ("--offset", "10", "offset", 10),
        ],
    )
    def test_optional_flags(
        self,
        flag: str,
        value: str,
        attr: str,
        expected: object,
    ) -> None:
        with patch(
            "sys.argv",
            ["basic_check", "--mode", "x", flag, value],
        ):
            args = parse_args()
        assert getattr(args, attr) == expected

    def test_defaults(self) -> None:
        with patch("sys.argv", ["basic_check", "--mode", "x"]):
            args = parse_args()
        assert args.num_keys == 100
        assert args.concurrency == 16
        assert args.offset == 0
        assert args.model == "/lmcache_test_model/"


# --------------- main() tests ---------------


class TestMain:
    """Tests for basic_check.main (async)."""

    @pytest.mark.asyncio
    async def test_list_mode(self, monkeypatch, capsys) -> None:
        """--mode list should print available modes."""
        monkeypatch.setattr("sys.argv", ["basic_check", "--mode", "list"])

        fake_reg = CheckModeRegistry()
        fake_reg.loaded = True
        fake_reg.modes["alpha"] = lambda: None
        fake_reg.modes["beta"] = lambda: None

        monkeypatch.setattr("lmcache.v1.basic_check.registry", fake_reg)

        await main()
        out = capsys.readouterr().out
        assert "alpha" in out
        assert "beta" in out

    @pytest.mark.asyncio
    async def test_unknown_mode(self, monkeypatch, capsys) -> None:
        """Unknown mode should print an error."""
        monkeypatch.setattr(
            "sys.argv",
            ["basic_check", "--mode", "no_such_mode"],
        )

        fake_reg = CheckModeRegistry()
        fake_reg.loaded = True  # empty registry

        monkeypatch.setattr("lmcache.v1.basic_check.registry", fake_reg)

        await main()
        out = capsys.readouterr().out
        assert "Unknown mode" in out

    @pytest.mark.asyncio
    async def test_valid_mode_invokes_function(self, monkeypatch) -> None:
        """A valid mode should be called with parsed args."""
        handler = AsyncMock()

        fake_reg = CheckModeRegistry()
        fake_reg.loaded = True
        fake_reg.register("run_it", handler)

        monkeypatch.setattr("lmcache.v1.basic_check.registry", fake_reg)
        monkeypatch.setattr(
            "sys.argv",
            [
                "basic_check",
                "--mode",
                "run_it",
                "--num-keys",
                "5",
                "--concurrency",
                "2",
                "--offset",
                "3",
            ],
        )

        await main()

        handler.assert_awaited_once_with(
            model="/lmcache_test_model/",
            num_keys=5,
            concurrency=2,
            offset=3,
        )
