# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Tests for the composable search steps and cascade runner."""

from __future__ import annotations

import os

import pytest

from cuda.pathfinder import UnsupportedArchError
from cuda.pathfinder._dynamic_libs import search_platform as search_platform_mod
from cuda.pathfinder._dynamic_libs.descriptor_catalog import WindowsSearchDirs
from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS, LibDescriptor
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError
from cuda.pathfinder._dynamic_libs.search_platform import LinuxSearchPlatform, WindowsSearchPlatform
from cuda.pathfinder._dynamic_libs.search_steps import (
    EARLY_FIND_STEPS,
    LATE_FIND_STEPS,
    FindResult,
    SearchContext,
    _find_lib_dir_using_anchor,
    find_in_conda,
    find_in_cuda_path,
    find_in_site_packages,
    run_find_steps,
)
from cuda.pathfinder._utils import windows_arch as windows_arch_mod

_STEPS_MOD = "cuda.pathfinder._dynamic_libs.search_steps"
_PLAT_MOD = "cuda.pathfinder._dynamic_libs.search_platform"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_desc(name: str = "cudart", **overrides) -> LibDescriptor:
    defaults = {
        "name": name,
        "packaged_with": "ctk",
        "linux_sonames": ("libcudart.so",),
        "windows_dlls": ("cudart64_12.dll",),
        "site_packages_linux": (os.path.join("nvidia", "cuda_runtime", "lib"),),
        "site_packages_windows": WindowsSearchDirs(
            x64=(os.path.join("nvidia", "cuda_runtime", "bin"),),
            arm64=(os.path.join("nvidia", "cuda_runtime", "bin"),),
        ),
    }
    defaults.update(overrides)
    return LibDescriptor(**defaults)


def _ctx(desc: LibDescriptor | None = None, *, platform=None) -> SearchContext:
    if platform is None:
        platform = LinuxSearchPlatform()
    return SearchContext(desc or _make_desc(), platform=platform)


def _patch_site_packages_search(mocker, root):
    def _find_sub_dirs(sub_dirs):
        path = root.joinpath(*sub_dirs)
        return [str(path)] if path.is_dir() else []

    return mocker.patch(f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", side_effect=_find_sub_dirs)


def _write_pe(path, machine):
    image = bytearray(0x86)
    image[:2] = b"MZ"
    image[0x3C:0x40] = (0x80).to_bytes(4, "little")
    image[0x80:0x84] = b"PE\0\0"
    image[0x84:0x86] = machine.to_bytes(2, "little")
    path.write_bytes(image)


# ---------------------------------------------------------------------------
# SearchContext
# ---------------------------------------------------------------------------


class TestSearchContext:
    def test_libname_delegates_to_descriptor(self):
        ctx = _ctx(_make_desc(name="nvrtc"))
        assert ctx.libname == "nvrtc"

    def test_lib_searched_for_linux(self):
        ctx = SearchContext(_make_desc(name="cublas"), platform=LinuxSearchPlatform())
        assert ctx.lib_searched_for == "libcublas.so"

    def test_lib_searched_for_windows(self):
        ctx = SearchContext(_make_desc(name="cublas"), platform=WindowsSearchPlatform(target_arch="x64"))
        assert ctx.lib_searched_for == "cublas*.dll"

    def test_raise_not_found_includes_messages(self):
        ctx = _ctx()
        ctx.error_messages.append("No such file: libcudart.so*")
        ctx.attachments.append('  listdir("/some/dir"):')
        with pytest.raises(DynamicLibNotFoundError, match="No such file"):
            ctx.raise_not_found()

    def test_raise_not_found_empty_messages(self):
        ctx = _ctx()
        with pytest.raises(DynamicLibNotFoundError):
            ctx.raise_not_found()


# ---------------------------------------------------------------------------
# Windows Python architecture detection
# ---------------------------------------------------------------------------


class TestWindowsPythonArch:
    @pytest.mark.agent_authored(model="gpt-5")
    def test_linux_platform_does_not_detect_windows_arch(self, mocker):
        mocker.patch.object(search_platform_mod, "IS_WINDOWS", False)
        get_windows_arch = mocker.patch.object(search_platform_mod, "windows_python_arch")

        platform = search_platform_mod._platform_for_current_system()

        assert isinstance(platform, LinuxSearchPlatform)
        get_windows_arch.assert_not_called()

    @pytest.mark.agent_authored(model="gpt-5")
    def test_detects_sysconfig_x64(self, mocker):
        mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="win-amd64")

        assert windows_arch_mod.windows_python_arch() == "x64"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_detects_sysconfig_arm64(self, mocker):
        mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="win-arm64")

        assert windows_arch_mod.windows_python_arch() == "arm64"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_rejects_unknown_sysconfig_tag(self, mocker):
        mocker.patch.object(windows_arch_mod.sysconfig, "get_platform", return_value="custom-win")

        with pytest.raises(
            UnsupportedArchError,
            match=r"Unsupported Windows Python platform tag: 'custom-win'.*win-amd64.*win-arm64",
        ) as exc_info:
            windows_arch_mod.windows_python_arch()
        assert exc_info.value.platform_tag == "custom-win"


@pytest.mark.parametrize(
    ("machine", "target_arch", "expected"),
    (
        (0x8664, "x64", True),
        (0x8664, "arm64", False),
        (0xAA64, "x64", False),
        (0xAA64, "arm64", True),
    ),
)
@pytest.mark.agent_authored(model="gpt-5")
def test_windows_pe_matches_arch(tmp_path, machine, target_arch, expected):
    dll = tmp_path / "test.dll"
    _write_pe(dll, machine)

    assert windows_arch_mod.windows_pe_matches_arch(str(dll), target_arch) is expected


@pytest.mark.agent_authored(model="gpt-5")
def test_windows_pe_matches_arch_rejects_malformed_file(tmp_path):
    dll = tmp_path / "test.dll"
    dll.write_bytes(b"not a PE file")

    assert windows_arch_mod.windows_pe_matches_arch(str(dll), "x64") is False


# ---------------------------------------------------------------------------
# find_in_site_packages
# ---------------------------------------------------------------------------


class TestFindInSitePackages:
    def test_returns_none_when_no_rel_dirs(self):
        desc = _make_desc(site_packages_linux=(), site_packages_windows=WindowsSearchDirs())
        result = find_in_site_packages(_ctx(desc))
        assert result is None

    def test_found_linux(self, mocker, tmp_path):
        lib_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
        lib_dir.mkdir(parents=True)
        so_file = lib_dir / "libcudart.so"
        so_file.touch()

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(lib_dir)],
        )

        desc = _make_desc(
            site_packages_linux=(os.path.join("nvidia", "cuda_runtime", "lib"),),
        )
        result = find_in_site_packages(_ctx(desc, platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(so_file)
        assert result.found_via == "site-packages"

    def test_found_windows(self, mocker, tmp_path):
        bin_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin"
        bin_dir.mkdir(parents=True)
        dll = bin_dir / "cudart64_12.dll"
        dll.touch()

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(bin_dir)],
        )
        mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False)

        desc = _make_desc(
            name="cudart",
            site_packages_windows=WindowsSearchDirs(
                x64=(os.path.join("nvidia", "cuda_runtime", "bin"),),
                arm64=(os.path.join("nvidia", "cuda_runtime", "bin"),),
            ),
        )
        result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="x64")))
        assert result is not None
        assert result.abs_path == str(dll)
        assert result.found_via == "site-packages"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_arm64_prefers_cuda13_arch_dir_to_cuda12(self, mocker, tmp_path):
        x86_64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "x86_64"
        arm64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "arm64"
        cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin"
        x86_64_dir.mkdir(parents=True)
        arm64_dir.mkdir(parents=True)
        cuda12_dir.mkdir(parents=True)
        (x86_64_dir / "cudart64_12.dll").touch()
        (cuda12_dir / "cudart64_12.dll").touch()
        arm64_dll = arm64_dir / "cudart64_12.dll"
        arm64_dll.touch()

        _patch_site_packages_search(mocker, tmp_path)
        mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False)

        desc = LIB_DESCRIPTORS["cudart"]
        result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="arm64")))
        assert result is not None
        assert result.abs_path == str(arm64_dll)
        assert result.found_via == "site-packages"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_x64_prefers_cuda13_arch_dir_to_cuda12(self, mocker, tmp_path):
        x86_64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "x86_64"
        arm64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "arm64"
        cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin"
        x86_64_dir.mkdir(parents=True)
        arm64_dir.mkdir(parents=True)
        cuda12_dir.mkdir(parents=True)
        x86_64_dll = x86_64_dir / "cudart64_12.dll"
        x86_64_dll.touch()
        (arm64_dir / "cudart64_12.dll").touch()
        (cuda12_dir / "cudart64_12.dll").touch()

        _patch_site_packages_search(mocker, tmp_path)
        mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False)

        desc = LIB_DESCRIPTORS["cudart"]
        result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="x64")))
        assert result is not None
        assert result.abs_path == str(x86_64_dll)
        assert result.found_via == "site-packages"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_x64_uses_cuda12_when_cuda13_is_absent(self, mocker, tmp_path):
        cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin"
        cuda12_dir.mkdir(parents=True)
        cuda12_dll = cuda12_dir / "cudart64_12.dll"
        cuda12_dll.touch()

        _patch_site_packages_search(mocker, tmp_path)
        mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False)

        desc = LIB_DESCRIPTORS["cudart"]
        result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform(target_arch="x64")))
        assert result is not None
        assert result.abs_path == str(cuda12_dll)
        assert result.found_via == "site-packages"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_arm64_skips_cuda12_when_cuda13_is_absent(self, mocker, tmp_path):
        cuda12_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin"
        cuda12_dir.mkdir(parents=True)
        (cuda12_dir / "cudart64_12.dll").touch()

        _patch_site_packages_search(mocker, tmp_path)
        mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False)

        desc = LIB_DESCRIPTORS["cudart"]
        platform = WindowsSearchPlatform(target_arch="arm64")
        assert platform.site_packages_rel_dirs(desc) == ("nvidia/cu13/bin/arm64",)

        result = find_in_site_packages(_ctx(desc, platform=platform))

        assert result is None

    def test_not_found_appends_error(self, mocker, tmp_path):
        empty_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
        empty_dir.mkdir(parents=True)

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(empty_dir)],
        )

        ctx = _ctx(platform=LinuxSearchPlatform())
        result = find_in_site_packages(ctx)
        assert result is None
        assert any("No such file" in m for m in ctx.error_messages)

    # The next three tests cover the Linux glob fallback in
    # cuda.pathfinder._dynamic_libs.search_platform._find_so_in_rel_dirs.
    # The fallback triggers when the unversioned libfoo.so is absent but
    # versioned libfoo.so.<major> files exist (e.g. some conda layouts).
    # Issue #1732 tracks the decision to return the newest-sorting match
    # deterministically; these tests lock in that policy at the
    # site-packages call site.

    def test_glob_fallback_returns_single_versioned_match(self, mocker, tmp_path):
        lib_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
        lib_dir.mkdir(parents=True)
        versioned = lib_dir / "libcudart.so.13"
        versioned.touch()

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(lib_dir)],
        )

        result = find_in_site_packages(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(versioned)
        assert result.found_via == "site-packages"

    def test_glob_fallback_returns_newest_of_multiple_matches(self, mocker, tmp_path):
        lib_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
        lib_dir.mkdir(parents=True)
        older = lib_dir / "libcudart.so.12"
        newer = lib_dir / "libcudart.so.13"
        older.touch()
        newer.touch()

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(lib_dir)],
        )

        result = find_in_site_packages(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(newer)
        assert result.found_via == "site-packages"

    def test_glob_fallback_zero_matches_returns_none(self, mocker, tmp_path):
        lib_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib"
        lib_dir.mkdir(parents=True)
        (lib_dir / "unrelated.txt").touch()

        mocker.patch(
            f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages",
            return_value=[str(lib_dir)],
        )

        ctx = _ctx(platform=LinuxSearchPlatform())
        result = find_in_site_packages(ctx)
        assert result is None
        assert any("No such file" in m and "libcudart.so" in m for m in ctx.error_messages)


# ---------------------------------------------------------------------------
# find_in_conda
# ---------------------------------------------------------------------------


class TestFindInConda:
    def test_returns_none_without_conda_prefix(self, mocker):
        mocker.patch.dict(os.environ, {}, clear=True)
        assert find_in_conda(_ctx()) is None

    def test_returns_none_with_empty_conda_prefix(self, mocker):
        mocker.patch.dict(os.environ, {"CONDA_PREFIX": ""})
        assert find_in_conda(_ctx()) is None

    def test_found_linux(self, mocker, tmp_path):
        lib_dir = tmp_path / "lib"
        lib_dir.mkdir()
        so_file = lib_dir / "libcudart.so"
        so_file.touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        result = find_in_conda(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(so_file)
        assert result.found_via == "conda"

    def test_found_windows(self, mocker, tmp_path):
        bin_dir = tmp_path / "Library" / "bin"
        bin_dir.mkdir(parents=True)
        dll = bin_dir / "cudart64_12.dll"
        dll.touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        result = find_in_conda(_ctx(platform=WindowsSearchPlatform(target_arch="x64")))
        assert result is not None
        assert result.abs_path == str(dll)
        assert result.found_via == "conda"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_arm64_prefers_arch_dir(self, mocker, tmp_path):
        x64_dir = tmp_path / "Library" / "bin" / "x64"
        arm64_dir = tmp_path / "Library" / "bin" / "arm64"
        x64_dir.mkdir(parents=True)
        arm64_dir.mkdir(parents=True)
        (x64_dir / "cudart64_12.dll").touch()
        arm64_dll = arm64_dir / "cudart64_12.dll"
        arm64_dll.touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        result = find_in_conda(_ctx(platform=WindowsSearchPlatform(target_arch="arm64")))
        assert result is not None
        assert result.abs_path == str(arm64_dll)
        assert result.found_via == "conda"

    # The next three tests cover the Linux glob fallback in
    # cuda.pathfinder._dynamic_libs.search_platform.LinuxSearchPlatform.find_in_lib_dir,
    # which is exercised by find_in_conda (and find_in_cuda_path) when the
    # resolved lib dir contains only versioned libfoo.so.<major> files.
    # Issue #1732 tracks the decision to return the newest-sorting match
    # deterministically; these tests lock in that policy at the conda /
    # CUDA_PATH call site.

    def test_glob_fallback_returns_single_versioned_match(self, mocker, tmp_path):
        lib_dir = tmp_path / "lib"
        lib_dir.mkdir()
        versioned = lib_dir / "libcudart.so.13"
        versioned.touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        result = find_in_conda(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(versioned)
        assert result.found_via == "conda"

    def test_glob_fallback_returns_newest_of_multiple_matches(self, mocker, tmp_path):
        lib_dir = tmp_path / "lib"
        lib_dir.mkdir()
        older = lib_dir / "libcudart.so.12"
        newer = lib_dir / "libcudart.so.13"
        older.touch()
        newer.touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        result = find_in_conda(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(newer)
        assert result.found_via == "conda"

    def test_glob_fallback_zero_matches_returns_none(self, mocker, tmp_path):
        lib_dir = tmp_path / "lib"
        lib_dir.mkdir()
        (lib_dir / "unrelated.txt").touch()

        mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)})

        ctx = _ctx(platform=LinuxSearchPlatform())
        result = find_in_conda(ctx)
        assert result is None
        assert any("No such file" in m and "libcudart.so" in m for m in ctx.error_messages)


# ---------------------------------------------------------------------------
# find_in_cuda_path
# ---------------------------------------------------------------------------


class TestFindInCudaHome:
    def test_returns_none_without_env_var(self, mocker):
        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=None)
        assert find_in_cuda_path(_ctx(platform=LinuxSearchPlatform())) is None

    def test_found_linux(self, mocker, tmp_path):
        lib_dir = tmp_path / "lib64"
        lib_dir.mkdir()
        so_file = lib_dir / "libcudart.so"
        so_file.touch()

        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path))

        result = find_in_cuda_path(_ctx(platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(so_file)
        assert result.found_via == "CUDA_PATH"

    def test_found_windows(self, mocker, tmp_path):
        bin_dir = tmp_path / "bin"
        bin_dir.mkdir()
        dll = bin_dir / "cudart64_12.dll"
        dll.touch()

        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path))

        result = find_in_cuda_path(_ctx(platform=WindowsSearchPlatform(target_arch="x64")))
        assert result is not None
        assert result.abs_path == str(dll)
        assert result.found_via == "CUDA_PATH"

    @pytest.mark.agent_authored(model="gpt-5")
    def test_found_windows_arm64_prefers_arch_dir(self, mocker, tmp_path):
        x64_dir = tmp_path / "bin" / "x64"
        arm64_dir = tmp_path / "bin" / "arm64"
        x64_dir.mkdir(parents=True)
        arm64_dir.mkdir(parents=True)
        (x64_dir / "cudart64_12.dll").touch()
        arm64_dll = arm64_dir / "cudart64_12.dll"
        arm64_dll.touch()

        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path))

        result = find_in_cuda_path(_ctx(platform=WindowsSearchPlatform(target_arch="arm64")))
        assert result is not None
        assert result.abs_path == str(arm64_dll)
        assert result.found_via == "CUDA_PATH"

    @pytest.mark.parametrize(
        ("target_arch", "machine", "expected_found"),
        (
            ("x64", 0x8664, True),
            ("x64", 0xAA64, False),
            ("arm64", 0x8664, False),
            ("arm64", 0xAA64, True),
        ),
    )
    @pytest.mark.agent_authored(model="gpt-5")
    def test_nvvm_windows_checks_binary_arch(self, mocker, tmp_path, target_arch, machine, expected_found):
        nvvm_dir = tmp_path / "nvvm" / "bin"
        nvvm_dir.mkdir(parents=True)
        dll = nvvm_dir / "nvvm64_40_0.dll"
        _write_pe(dll, machine)

        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path))

        ctx = _ctx(LIB_DESCRIPTORS["nvvm"], platform=WindowsSearchPlatform(target_arch=target_arch))
        result = find_in_cuda_path(ctx)

        assert (result is not None) is expected_found
        if expected_found:
            assert result is not None
            assert result.abs_path == str(dll)
            assert result.found_via == "CUDA_PATH"
        else:
            assert any(f"No {target_arch}-compatible PE file" in message for message in ctx.error_messages)


# ---------------------------------------------------------------------------
# run_find_steps
# ---------------------------------------------------------------------------


class TestRunFindSteps:
    def test_returns_first_hit(self):
        hit = FindResult("/path/to/lib.so", "step-a")

        def step_a(_ctx):
            return hit

        def step_b(_ctx):
            raise AssertionError("step_b should not be called")

        result = run_find_steps(_ctx(), (step_a, step_b))
        assert result is hit

    def test_returns_none_when_all_miss(self):
        result = run_find_steps(_ctx(), (lambda _: None, lambda _: None))
        assert result is None

    def test_empty_steps(self):
        assert run_find_steps(_ctx(), ()) is None

    def test_skips_nones_returns_later_hit(self):
        hit = FindResult("/later/lib.so", "step-c")
        result = run_find_steps(_ctx(), (lambda _: None, lambda _: hit))
        assert result is hit


# ---------------------------------------------------------------------------
# Step tuple sanity checks
# ---------------------------------------------------------------------------


class TestStepTuples:
    def test_early_find_steps_contains_expected(self):
        assert find_in_site_packages in EARLY_FIND_STEPS
        assert find_in_conda in EARLY_FIND_STEPS

    def test_late_find_steps_contains_expected(self):
        assert find_in_cuda_path in LATE_FIND_STEPS

    def test_early_and_late_are_disjoint(self):
        assert not set(EARLY_FIND_STEPS) & set(LATE_FIND_STEPS)


# ---------------------------------------------------------------------------
# Data-driven anchor paths
# ---------------------------------------------------------------------------


class TestAnchorRelDirs:
    """Verify that descriptor anchor paths drive directory resolution."""

    @pytest.mark.agent_authored(model="gpt-5")
    def test_windows_search_dirs_arch_only_constructors(self):
        assert WindowsSearchDirs.x64_only("first", "second") == WindowsSearchDirs(x64=("first", "second"))
        assert WindowsSearchDirs.arm64_only("first", "second") == WindowsSearchDirs(arm64=("first", "second"))

    def test_nvvm_has_custom_linux_paths(self):
        desc = LIB_DESCRIPTORS["nvvm"]
        assert desc.anchor_rel_dirs_linux == ("nvvm/lib64",)

    def test_nvvm_has_custom_windows_paths(self):
        desc = LIB_DESCRIPTORS["nvvm"]
        assert desc.anchor_rel_dirs_windows.for_arch("x64") == ("nvvm/bin/x64", "nvvm/bin")
        assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("nvvm/bin",)

    @pytest.mark.agent_authored(model="gpt-5")
    def test_cupti_has_custom_windows_paths(self):
        desc = LIB_DESCRIPTORS["cupti"]
        assert desc.anchor_rel_dirs_windows.for_arch("x64") == (
            "extras/CUPTI/lib/x64",
            "extras/CUPTI/lib64",
            "bin",
        )
        assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("extras/CUPTI/lib/arm64",)

    @pytest.mark.parametrize("libname", ["cudart", "cublas", "nvrtc"])
    def test_regular_ctk_libs_use_defaults(self, libname):
        desc = LIB_DESCRIPTORS[libname]
        assert desc.anchor_rel_dirs_linux == ("lib64", "lib")
        assert desc.anchor_rel_dirs_windows.for_arch("x64") == ("bin/x64", "bin")
        assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("bin/arm64",)

    @pytest.mark.agent_authored(model="gpt-5")
    def test_cudla_uses_arm64_only_windows_anchor(self):
        desc = LIB_DESCRIPTORS["cudla"]

        assert desc.anchor_rel_dirs_windows.for_arch("x64") == ()
        assert desc.anchor_rel_dirs_windows.for_arch("arm64") == ("bin/arm64",)

    @pytest.mark.agent_authored(model="gpt-5")
    def test_windows_anchor_dirs_select_arm64(self):
        desc = _make_desc(
            anchor_rel_dirs_windows=WindowsSearchDirs(
                x64=("bin/x64", "bin"),
                arm64=("bin/arm64", "bin"),
            )
        )
        assert WindowsSearchPlatform(target_arch="arm64").anchor_rel_dirs(desc) == ("bin/arm64", "bin")

    @pytest.mark.agent_authored(model="gpt-5")
    def test_windows_anchor_dirs_select_x64(self):
        desc = _make_desc(
            anchor_rel_dirs_windows=WindowsSearchDirs(
                x64=("bin/x64", "bin"),
                arm64=("bin/arm64", "bin"),
            )
        )
        assert WindowsSearchPlatform(target_arch="x64").anchor_rel_dirs(desc) == ("bin/x64", "bin")

    def test_find_lib_dir_uses_descriptor_linux(self, tmp_path):
        (tmp_path / "nvvm" / "lib64").mkdir(parents=True)

        desc = _make_desc(name="nvvm", anchor_rel_dirs_linux=("nvvm/lib64",))
        result = _find_lib_dir_using_anchor(desc, LinuxSearchPlatform(), str(tmp_path))
        assert result is not None
        assert result.endswith(os.path.join("nvvm", "lib64"))

    def test_find_lib_dir_uses_descriptor_windows(self, tmp_path):
        (tmp_path / "nvvm" / "bin").mkdir(parents=True)

        desc = _make_desc(
            name="nvvm",
            anchor_rel_dirs_windows=WindowsSearchDirs(
                x64=("nvvm/bin/x64", "nvvm/bin"),
                arm64=("nvvm/bin/arm64", "nvvm/bin"),
            ),
        )
        result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(target_arch="x64"), str(tmp_path))
        assert result is not None
        assert result.endswith(os.path.join("nvvm", "bin"))

    @pytest.mark.agent_authored(model="gpt-5")
    def test_find_lib_dir_windows_arm64_uses_arm64_anchor(self, tmp_path):
        (tmp_path / "bin" / "x64").mkdir(parents=True)
        (tmp_path / "bin" / "arm64").mkdir(parents=True)

        desc = _make_desc(
            name="cudart",
            anchor_rel_dirs_windows=WindowsSearchDirs(
                x64=("bin/x64", "bin"),
                arm64=("bin/arm64",),
            ),
        )
        result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(target_arch="arm64"), str(tmp_path))
        assert result is not None
        assert result.endswith(os.path.join("bin", "arm64"))

    def test_find_lib_dir_returns_none_when_no_match(self, tmp_path):
        desc = _make_desc(anchor_rel_dirs_linux=("nonexistent",))
        assert _find_lib_dir_using_anchor(desc, LinuxSearchPlatform(), str(tmp_path)) is None

    def test_nvvm_cuda_home_linux(self, mocker, tmp_path):
        """End-to-end: find_in_cuda_path resolves nvvm under its custom subdir."""
        mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(tmp_path))

        nvvm_dir = tmp_path / "nvvm" / "lib64"
        nvvm_dir.mkdir(parents=True)
        so_file = nvvm_dir / "libnvvm.so"
        so_file.touch()

        desc = _make_desc(
            name="nvvm",
            linux_sonames=("libnvvm.so",),
            anchor_rel_dirs_linux=("nvvm/lib64",),
        )
        result = find_in_cuda_path(_ctx(desc, platform=LinuxSearchPlatform()))
        assert result is not None
        assert result.abs_path == str(so_file)
        assert result.found_via == "CUDA_PATH"
