"""
Test that the CLI commands work as expected.

In general there can be two types of tests for each command:
- Real tests (with suffix `_real`) that invoke the commands without any mocking
- Mocked tests (with suffix `_mocked`) that use monkeypatch to mock logic that would
  otherwise be slow (e.g. downloading cubins, filesystem calls, etc), and also to
  create deterministic state so we can check for expected output (e.g. number of cubins)

These tests don't require a GPU. CLI tests that require a GPU are in test_cli_cmds_gpu.py.

Note: The `replay` command is tested in tests/utils/test_logging_replay.py alongside
the other logging/replay functionality tests, since it's tightly coupled with that feature.
"""

from .cli_cmd_helpers import (
    _test_cmd_helper,
    _assert_output_contains_all,
    _assert_output_contains_any,
)
from flashinfer.artifacts import ArtifactPath
from packaging.version import Version
import sys


def test_show_config_cmd_real():
    """
    Test that show-config command works as expected
    """
    out = _test_cmd_helper(["show-config"])

    # Basic sections present
    _assert_output_contains_all(
        out,
        "=== Torch Version Info ===",
        "=== Environment Variables ===",
        "=== Artifact Path ===",
        "=== Downloaded Cubins ===",
    )


def test_show_config_cmd_mocked(monkeypatch):
    """
    Test that show-config command works as but with mocked cubin status
    """
    # Don't check filesystem for cubins
    monkeypatch.setattr(
        "flashinfer.__main__.get_artifacts_status",
        lambda: (("foo.cubin", True), ("bar.cubin", False)),
    )
    # Avoid module registration/inspection
    monkeypatch.setattr(
        "flashinfer.__main__._ensure_modules_registered",
        lambda: [],
    )

    out = _test_cmd_helper(["show-config"])

    # Uses our monkeypatched data
    assert "Downloaded 1/2 cubins" in out


def test_cli_group_help_real():
    """
    Test that the CLI group runs without error and sanity checks the output
    """
    out = _test_cmd_helper([])
    _assert_output_contains_any(out, "FlashInfer CLI", "Usage")


def test_download_cubin_flag_mocked(monkeypatch):
    # This just tests that the flag is parsed correctly, so we can monkeypatch
    # download_artifacts to avoid the latency of downloading cubins
    monkeypatch.setattr("flashinfer.__main__.download_artifacts", lambda: None)

    out = _test_cmd_helper(["--download-cubin"])
    assert "All cubin download tasks completed successfully" in out


def test_download_cubin_cmd_mocked(monkeypatch):
    """
    Test that download-cubin downloads raw cubin artifacts using a mocked cubin path.
    """
    fmha_cubin = "fmhaSm100aKernel_QE4m3KvE2m1OE4m3H128PagedKvCausalP32VarSeqQ128Kv128PersistentContext.cubin"

    def mock_get_subdir_file_list():
        return [(f"{ArtifactPath.TRTLLM_GEN_FMHA}/{fmha_cubin}", "fake_checksum_12345")]

    monkeypatch.setattr(
        "flashinfer.artifacts.get_subdir_file_list", mock_get_subdir_file_list
    )
    monkeypatch.setattr(
        "flashinfer.artifacts.download_file", lambda *_args, **_kwargs: True
    )
    monkeypatch.setattr("flashinfer.artifacts.verify_cubin", lambda *_args: True)

    out = _test_cmd_helper(["download-cubin"])
    assert "All cubin download tasks completed successfully" in out


def test_install_cubin_wheel_cmd_mocked(monkeypatch):
    recorded = {}

    def mock_run(cmd, check=False):
        recorded["cmd"] = cmd
        recorded["check"] = check

        class Result:
            returncode = 0

        return Result()

    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1")
    monkeypatch.setattr("flashinfer.__main__.subprocess.run", mock_run)

    out = _test_cmd_helper(["install-cubin-wheel"])

    _assert_output_contains_all(
        out,
        "=== Cubin Wheel Install ===",
        "FlashInfer version:",
        "Wheel index:",
        "Requirement:",
        "flashinfer-cubin==0.4.1",
    )
    assert recorded["cmd"] == [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--upgrade",
        "--no-deps",
        "--index-url",
        "https://flashinfer.ai/whl",
        "flashinfer-cubin==0.4.1",
    ]
    assert recorded["check"] is False


def test_install_cubin_wheel_cmd_nightly_dry_run(monkeypatch):
    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1.dev20260421")

    def fail_run(*_args, **_kwargs):
        raise AssertionError("subprocess.run should not be called for dry-run")

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", fail_run)

    out = _test_cmd_helper(["install-cubin-wheel", "--nightly", "--dry-run"])

    _assert_output_contains_all(
        out,
        "https://flashinfer.ai/whl/nightly",
        "flashinfer-cubin==0.4.1.dev20260421",
        "Dry run requested; pip install was not executed.",
    )


def test_list_cubins_cmd_real():
    out = _test_cmd_helper(["list-cubins"])

    _assert_output_contains_all(out, "Cubin", "Status")


def test_list_cubins_cmd_mocked(monkeypatch):
    monkeypatch.setattr(
        "flashinfer.__main__.get_artifacts_status",
        lambda: (("foo.cubin", True), ("bar.cubin", False)),
    )

    out = _test_cmd_helper(["list-cubins"])
    _assert_output_contains_all(out, "foo.cubin", "bar.cubin")


def test_clear_cache_cmd_mocked(monkeypatch):
    """
    Test that clear-cache command works without actually clearing the cache.

    This doesn't test much, just a basic sanity check.
    """
    monkeypatch.setattr("flashinfer.__main__.clear_cache_dir", lambda: None)

    out = _test_cmd_helper(["clear-cache"])
    assert "Cache cleared successfully" in out


def test_clear_cache_cmd_real(monkeypatch, tmp_path):
    """
    Test that clear-cache command actually clears the cache directory.

    Uses a temporary directory to avoid side effects on the real cache.
    """
    # Create a temporary JIT directory with some dummy cache files
    temp_jit_dir = tmp_path / "cached_ops"
    temp_jit_dir.mkdir(parents=True, exist_ok=True)

    # Create some dummy cached files to simulate a real cache
    dummy_module_dir = temp_jit_dir / "test_module_abc123"
    dummy_module_dir.mkdir(parents=True, exist_ok=True)
    (dummy_module_dir / "test_module.so").write_text("dummy shared library")
    (dummy_module_dir / "build.ninja").write_text("dummy build file")

    # Monkeypatch the FLASHINFER_JIT_DIR to point to our temp directory
    monkeypatch.setattr("flashinfer.jit.core.jit_env.FLASHINFER_JIT_DIR", temp_jit_dir)

    # Verify the cache directory exists before clearing
    assert temp_jit_dir.exists()
    assert (dummy_module_dir / "test_module.so").exists()

    # Run the clear-cache command
    out = _test_cmd_helper(["clear-cache"])
    assert "Cache cleared successfully" in out

    # Verify the cache directory has been removed
    assert not temp_jit_dir.exists()


def test_clear_cubin_cmd_mocked(monkeypatch):
    """
    Test that clear-cubin command works without actually clearing the cubin.

    This doesn't test much, just a basic sanity check.
    """
    monkeypatch.setattr("flashinfer.__main__.clear_cubin", lambda: None)

    out = _test_cmd_helper(["clear-cubin"])
    assert "Cubin cleared successfully" in out


def test_clear_cubin_cmd_real(monkeypatch, tmp_path):
    """
    Test that clear-cubin command actually clears the cubin directory.

    Uses a temporary directory to avoid side effects on the real cubins.
    """
    # Create a temporary cubin directory with some dummy cubin files
    temp_cubin_dir = tmp_path / "cubins"
    temp_cubin_dir.mkdir(parents=True, exist_ok=True)

    # Create some dummy cubin files to simulate real cubins
    dummy_cubin_subdir = temp_cubin_dir / "trtllm_gen_fmha"
    dummy_cubin_subdir.mkdir(parents=True, exist_ok=True)
    (dummy_cubin_subdir / "test_kernel.cubin").write_text("dummy cubin data")
    (dummy_cubin_subdir / "checksums.txt").write_text("abc123 test_kernel.cubin")

    # Monkeypatch FLASHINFER_CUBIN_DIR to point to our temp directory
    # Need to patch it in multiple places where it's imported
    monkeypatch.setattr("flashinfer.artifacts.FLASHINFER_CUBIN_DIR", temp_cubin_dir)
    monkeypatch.setattr(
        "flashinfer.jit.cubin_loader.FLASHINFER_CUBIN_DIR", temp_cubin_dir
    )

    # Verify the cubin directory exists before clearing
    assert temp_cubin_dir.exists()
    assert (dummy_cubin_subdir / "test_kernel.cubin").exists()

    # Run the clear-cubin command
    out = _test_cmd_helper(["clear-cubin"])
    assert "Cubin cleared successfully" in out

    # Verify the cubin directory has been removed
    assert not temp_cubin_dir.exists()


def test_install_jit_cache_wheel_cmd_mocked(monkeypatch):
    import flashinfer.__main__ as flashinfer_main

    recorded = {}

    def mock_run(cmd, check=False):
        recorded["cmd"] = cmd
        recorded["check"] = check

        class Result:
            returncode = 0

        return Result()

    def fail_get_cuda_version():
        raise AssertionError("get_cuda_version should not be called")

    monkeypatch.setattr(flashinfer_main.torch.version, "cuda", "12.9")
    monkeypatch.setattr("flashinfer.__main__.get_cuda_version", fail_get_cuda_version)
    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1")
    monkeypatch.setattr("flashinfer.__main__.subprocess.run", mock_run)

    out = _test_cmd_helper(["install-jit-cache-wheel"])

    _assert_output_contains_all(
        out,
        "=== JIT Cache Wheel Install ===",
        "FlashInfer version:",
        "CUDA version:",
        "Wheel index:",
        "Requirement:",
    )
    assert recorded["cmd"] == [
        sys.executable,
        "-m",
        "pip",
        "install",
        "--upgrade",
        "--no-deps",
        "--index-url",
        "https://flashinfer.ai/whl/cu129",
        "flashinfer-jit-cache==0.4.1+cu129",
    ]
    assert recorded["check"] is False


def test_install_jit_cache_wheel_cmd_falls_back_without_torch_cuda(monkeypatch):
    import flashinfer.__main__ as flashinfer_main

    monkeypatch.setattr(flashinfer_main.torch.version, "cuda", None)
    monkeypatch.setattr("flashinfer.__main__.get_cuda_version", lambda: Version("12.9"))
    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1")

    def fail_run(*_args, **_kwargs):
        raise AssertionError("subprocess.run should not be called for dry-run")

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", fail_run)

    out = _test_cmd_helper(["install-jit-cache-wheel", "--dry-run"])

    _assert_output_contains_all(
        out,
        "CUDA version: 12.9",
        "Wheel CUDA label: cu129",
        "flashinfer-jit-cache==0.4.1+cu129",
    )


def test_install_jit_cache_wheel_cmd_nightly_dry_run(monkeypatch):
    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1.dev20260421")

    def fail_run(*_args, **_kwargs):
        raise AssertionError("subprocess.run should not be called for dry-run")

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", fail_run)

    out = _test_cmd_helper(
        [
            "install-jit-cache-wheel",
            "--cuda-version",
            "cu130",
            "--nightly",
            "--dry-run",
        ]
    )

    _assert_output_contains_all(
        out,
        "https://flashinfer.ai/whl/nightly/cu130",
        "flashinfer-jit-cache==0.4.1.dev20260421+cu130",
        "Dry run requested; pip install was not executed.",
    )


def test_install_jit_cache_wheel_cmd_floors_newer_cuda_minor(monkeypatch):
    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1")

    def fail_run(*_args, **_kwargs):
        raise AssertionError("subprocess.run should not be called for dry-run")

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", fail_run)

    out = _test_cmd_helper(
        [
            "install-jit-cache-wheel",
            "--cuda-version",
            "13.3",
            "--dry-run",
        ]
    )

    _assert_output_contains_all(
        out,
        "CUDA version: 13.3",
        "Wheel CUDA label: cu130",
        "https://flashinfer.ai/whl/cu130",
        "flashinfer-jit-cache==0.4.1+cu130",
    )


def test_install_jit_cache_wheel_cmd_rejects_unsupported_cuda_floor():
    from click.testing import CliRunner

    from flashinfer.__main__ import cli

    runner = CliRunner()
    result = runner.invoke(
        cli,
        [
            "install-jit-cache-wheel",
            "--cuda-version",
            "12.7",
            "--dry-run",
        ],
    )

    assert result.exit_code != 0
    _assert_output_contains_all(
        result.output,
        "No compatible flashinfer-jit-cache wheel found for CUDA 12.7",
        "Supported wheel labels: cu128, cu129, cu130",
    )


def test_download_jit_cache_alias_cmd_mocked(monkeypatch):
    recorded = {}

    def mock_run(cmd, check=False):
        recorded["cmd"] = cmd

        class Result:
            returncode = 0

        return Result()

    monkeypatch.setattr("flashinfer.__main__.__version__", "0.4.1")
    monkeypatch.setattr("flashinfer.__main__.subprocess.run", mock_run)

    out = _test_cmd_helper(
        [
            "download-jit-cache",
            "--cuda-version",
            "12.8",
            "--flashinfer-version",
            "0.4.2",
        ]
    )

    _assert_output_contains_all(
        out,
        "https://flashinfer.ai/whl/cu128",
        "flashinfer-jit-cache==0.4.2+cu128",
    )
    assert recorded["cmd"][-1] == "flashinfer-jit-cache==0.4.2+cu128"


def test_download_kernels_cmd_mocked(monkeypatch):
    recorded = []

    def mock_run(cmd, check=False):
        recorded.append((cmd, check))

        class Result:
            returncode = 0

        return Result()

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", mock_run)

    out = _test_cmd_helper(
        ["download-kernels", "--cuda-version", "12.9", "--flashinfer-version", "0.4.1"]
    )

    _assert_output_contains_all(
        out,
        "=== Cubin Wheel Install ===",
        "=== JIT Cache Wheel Install ===",
        "flashinfer-cubin==0.4.1",
        "flashinfer-jit-cache==0.4.1+cu129",
    )
    assert recorded == [
        (
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--upgrade",
                "--no-deps",
                "--index-url",
                "https://flashinfer.ai/whl",
                "flashinfer-cubin==0.4.1",
            ],
            False,
        ),
        (
            [
                sys.executable,
                "-m",
                "pip",
                "install",
                "--upgrade",
                "--no-deps",
                "--index-url",
                "https://flashinfer.ai/whl/cu129",
                "flashinfer-jit-cache==0.4.1+cu129",
            ],
            False,
        ),
    ]


def test_download_kernels_cmd_allows_local_flashinfer_version(monkeypatch):
    def fail_run(*_args, **_kwargs):
        raise AssertionError("subprocess.run should not be called for dry-run")

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", fail_run)

    out = _test_cmd_helper(
        [
            "download-kernels",
            "--cuda-version",
            "12.9",
            "--flashinfer-version",
            "0.4.1+cu124",
            "--dry-run",
        ]
    )

    _assert_output_contains_all(
        out,
        "FlashInfer version: 0.4.1+cu124",
        "flashinfer-cubin==0.4.1",
        "flashinfer-jit-cache==0.4.1+cu129",
        "Dry run requested; pip install was not executed.",
    )


def test_download_kernels_cmd_attempts_jit_cache_after_cubin_failure(monkeypatch):
    from click.testing import CliRunner

    from flashinfer.__main__ import cli

    recorded = []

    def mock_run(cmd, check=False):
        recorded.append((cmd, check))

        class Result:
            returncode = 1 if cmd[-1].startswith("flashinfer-cubin==") else 0

        return Result()

    monkeypatch.setattr("flashinfer.__main__.subprocess.run", mock_run)

    runner = CliRunner()
    result = runner.invoke(
        cli,
        ["download-kernels", "--cuda-version", "12.9", "--flashinfer-version", "0.4.1"],
    )

    assert result.exit_code != 0
    _assert_output_contains_all(
        result.output,
        "flashinfer-cubin install failed: pip install failed with exit code 1.",
        "=== JIT Cache Wheel Install ===",
        "flashinfer-jit-cache==0.4.1+cu129",
        "flashinfer-jit-cache installed successfully",
        "One or more kernel wheel installs failed:",
    )
    assert len(recorded) == 2
    assert recorded[0][0][-1] == "flashinfer-cubin==0.4.1"
    assert recorded[1][0][-1] == "flashinfer-jit-cache==0.4.1+cu129"


class MockJitSpec:
    """Mock JitSpec for testing export-compile-commands."""

    def __init__(self, name, compile_commands):
        self.name = name
        self._compile_commands = compile_commands

    def get_compile_commands(self):
        return self._compile_commands


def test_export_compile_commands_mocked(monkeypatch, tmp_path):
    """
    Test that export-compile-commands writes correct JSON output.
    """
    # Create mock specs with compile commands
    mock_specs = {
        "module_a": MockJitSpec(
            "module_a",
            [
                {
                    "directory": "/path/to/build",
                    "command": "nvcc -c kernel_a.cu",
                    "file": "kernel_a.cu",
                }
            ],
        ),
        "module_b": MockJitSpec(
            "module_b",
            [
                {
                    "directory": "/path/to/build",
                    "command": "nvcc -c kernel_b.cu",
                    "file": "kernel_b.cu",
                }
            ],
        ),
    }

    monkeypatch.setattr("flashinfer.__main__._ensure_modules_registered", lambda: [])
    monkeypatch.setattr(
        "flashinfer.__main__.jit_spec_registry.get_all_specs", lambda: mock_specs
    )

    # Use tmp_path to write output file
    output_file = tmp_path / "compile_commands.json"
    out = _test_cmd_helper(["export-compile-commands", str(output_file)])

    assert "Successfully exported 2 compile commands" in out
    assert output_file.exists()

    # Verify JSON content
    import json

    with open(output_file) as f:
        commands = json.load(f)

    assert len(commands) == 2
    assert commands[0]["file"] == "kernel_a.cu"
    assert commands[1]["file"] == "kernel_b.cu"


def test_export_compile_commands_output_option(monkeypatch, tmp_path):
    """
    Test that --output option overrides PATH argument.
    """
    mock_specs = {
        "module_a": MockJitSpec(
            "module_a",
            [{"directory": "/build", "command": "nvcc -c a.cu", "file": "a.cu"}],
        ),
    }

    monkeypatch.setattr("flashinfer.__main__._ensure_modules_registered", lambda: [])
    monkeypatch.setattr(
        "flashinfer.__main__.jit_spec_registry.get_all_specs", lambda: mock_specs
    )

    # PATH argument should be ignored when --output is specified
    output_file = tmp_path / "custom_output.json"
    ignored_file = tmp_path / "ignored.json"
    out = _test_cmd_helper(
        ["export-compile-commands", str(ignored_file), "--output", str(output_file)]
    )

    assert "Successfully exported 1 compile commands" in out
    assert output_file.exists()
    assert not ignored_file.exists()


def test_export_compile_commands_no_modules(monkeypatch, tmp_path):
    """
    Test that export-compile-commands handles empty module registry.
    """
    monkeypatch.setattr("flashinfer.__main__._ensure_modules_registered", lambda: [])
    monkeypatch.setattr(
        "flashinfer.__main__.jit_spec_registry.get_all_specs", lambda: {}
    )

    output_file = tmp_path / "compile_commands.json"
    out = _test_cmd_helper(["export-compile-commands", str(output_file)])

    assert "No modules found" in out
    # File should not be created when no modules exist
    assert not output_file.exists()
