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

import inspect

import pytest

from cuda.core import Device, Linker, LinkerOptions, Program, ProgramOptions, _linker
from cuda.core._module import ObjectCode
from cuda.core._program import _can_load_generated_ptx
from cuda.core._utils.cuda_utils import CUDAError

ARCH = "sm_" + "".join(f"{i}" for i in Device().compute_capability)

kernel_a = """
extern __device__ int B();
extern __device__ int C(int a, int b);
__global__ void A() { int result = C(B(), 1);}
"""
device_function_b = "__device__ int B() { return 0; }"
device_function_c = "__device__ int C(int a, int b) { return a + b; }"

is_culink_backend = _linker._decide_nvjitlink_or_driver()
if not is_culink_backend:
    from cuda.bindings import nvjitlink

    nvJitLinkError = nvjitlink.nvJitLinkError
else:

    class nvJitLinkError(Exception):
        pass


@pytest.fixture
def compile_ptx_functions(init_cuda):
    # Without -rdc (relocatable device code) option, the generated ptx will not included any unreferenced
    # device functions, causing the link to fail
    object_code_a_ptx = Program(kernel_a, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    object_code_b_ptx = Program(device_function_b, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    object_code_c_ptx = Program(device_function_c, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")

    return object_code_a_ptx, object_code_b_ptx, object_code_c_ptx


@pytest.fixture
def compile_ltoir_functions(init_cuda):
    object_code_a_ltoir = Program(kernel_a, "c++", ProgramOptions(link_time_optimization=True)).compile("ltoir")
    object_code_b_ltoir = Program(device_function_b, "c++", ProgramOptions(link_time_optimization=True)).compile(
        "ltoir"
    )
    object_code_c_ltoir = Program(device_function_c, "c++", ProgramOptions(link_time_optimization=True)).compile(
        "ltoir"
    )

    return object_code_a_ltoir, object_code_b_ltoir, object_code_c_ltoir


options = [
    LinkerOptions(),
    LinkerOptions(arch=ARCH, verbose=True),
    LinkerOptions(arch=ARCH, max_register_count=32),
    LinkerOptions(arch=ARCH, optimization_level=3),
    LinkerOptions(arch=ARCH, debug=True),
    LinkerOptions(arch=ARCH, lineinfo=True),
]
if not is_culink_backend:
    options += [
        LinkerOptions(arch=ARCH, time=True),
        LinkerOptions(arch=ARCH, optimize_unused_variables=True),
        LinkerOptions(arch=ARCH, ptxas_options="-v"),
        LinkerOptions(arch=ARCH, ptxas_options=["-v", "--verbose"]),
        LinkerOptions(arch=ARCH, ptxas_options=("-v", "--verbose")),
        LinkerOptions(arch=ARCH, split_compile=0),
        LinkerOptions(arch=ARCH, split_compile_extended=1),
        # The following options are supported by nvjitlink and deprecated by culink
        LinkerOptions(arch=ARCH, ftz=True),
        LinkerOptions(arch=ARCH, prec_div=True),
        LinkerOptions(arch=ARCH, prec_sqrt=True),
        LinkerOptions(arch=ARCH, fma=True),
        LinkerOptions(arch=ARCH, kernels_used="A"),
        LinkerOptions(arch=ARCH, kernels_used=["C", "B"]),
        LinkerOptions(arch=ARCH, kernels_used=("C", "B")),
        LinkerOptions(arch=ARCH, variables_used="var1"),
        LinkerOptions(arch=ARCH, variables_used=["var1", "var2"]),
        LinkerOptions(arch=ARCH, variables_used=("var1", "var2")),
    ]
    version = nvjitlink.version()
    if version >= (12, 5):
        options.append(LinkerOptions(arch=ARCH, no_cache=True))


@pytest.mark.parametrize("options", options)
def test_linker_init(compile_ptx_functions, options):
    linker = Linker(*compile_ptx_functions, options=options)
    object_code = linker.link("cubin")
    assert isinstance(object_code, ObjectCode)
    assert Linker.which_backend() == ("driver" if is_culink_backend else "nvJitLink")


def test_linker_init_invalid_arch(compile_ptx_functions):
    err = AttributeError if is_culink_backend else nvjitlink.nvJitLinkError
    with pytest.raises(err):
        options = LinkerOptions(arch="99", ptx=True)
        Linker(*compile_ptx_functions, options=options)


@pytest.mark.skipif(is_culink_backend, reason="culink does not support ptx option")
def test_linker_link_ptx_nvjitlink(compile_ltoir_functions):
    options = LinkerOptions(arch=ARCH, link_time_optimization=True, ptx=True)
    linker = Linker(*compile_ltoir_functions, options=options)
    linked_code = linker.link("ptx")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


@pytest.mark.skipif(not is_culink_backend, reason="nvjitlink requires lto for ptx linking")
def test_linker_link_ptx_culink(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linked_code = linker.link("ptx")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_cubin(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linked_code = linker.link("cubin")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_ptx_multiple(compile_ptx_functions):
    ptxes = tuple(ObjectCode.from_ptx(obj.code) for obj in compile_ptx_functions)
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*ptxes, options=options)
    linked_code = linker.link("cubin")
    assert isinstance(linked_code, ObjectCode)
    assert linked_code.name == options.name


def test_linker_link_invalid_target_type(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    with pytest.raises(ValueError):
        linker.link("invalid_target")


def test_linker_get_error_log(compile_ptx_functions):
    options = LinkerOptions(name="ABC", arch=ARCH)

    replacement_kernel = """
extern __device__ int Z();
extern __device__ int C(int a, int b);
__global__ void A() { int result = C(Z(), 1);}
"""
    dummy_program = Program(
        replacement_kernel, "c++", ProgramOptions(name="CBA", relocatable_device_code=True)
    ).compile("ptx")
    linker = Linker(dummy_program, *(compile_ptx_functions[1:]), options=options)
    try:
        linker.link("cubin")

    except (nvJitLinkError, CUDAError):
        log = linker.get_error_log()
        assert isinstance(log, str)
        # TODO when 4902246 is addressed, we can update this to cover nvjitlink as well
        # The error is coming from the input object that's being linked (CBA), not the output object (ABC).
        if is_culink_backend:
            assert log.rstrip("\x00") == "error   : Undefined reference to '_Z1Zv' in 'CBA'"


def test_linker_get_info_log(compile_ptx_functions):
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linker.link("cubin")
    log = linker.get_info_log()
    assert isinstance(log, str)


@pytest.mark.skipif(is_culink_backend, reason="as_bytes() only supported for nvjitlink backend")
def test_linker_options_as_bytes_nvjitlink():
    """Test LinkerOptions.as_bytes() for nvJitLink backend"""
    options = LinkerOptions(arch="sm_80", debug=True, ftz=True, max_register_count=32)
    nvjitlink_options = options.as_bytes("nvjitlink")

    # Should return list of bytes
    assert isinstance(nvjitlink_options, list)
    assert all(isinstance(opt, bytes) for opt in nvjitlink_options)

    # Decode to check content
    options_str = [opt.decode() for opt in nvjitlink_options]
    assert "-arch=sm_80" in options_str
    assert "-g" in options_str
    assert "-ftz=true" in options_str
    assert "-maxrregcount=32" in options_str


@pytest.mark.parametrize("backend", ("invalid", "driver"))
def test_linker_options_as_bytes_invalid_backend(backend):
    """Test LinkerOptions.as_bytes() with invalid backend"""
    options = LinkerOptions(arch="sm_80")
    with pytest.raises(ValueError, match="only supports 'nvjitlink' backend"):
        options.as_bytes(backend)


def test_linker_logs_cached_after_link(compile_ptx_functions):
    """After a successful link(), get_error_log/get_info_log should return cached strings."""
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    linker.link("cubin")
    err_log = linker.get_error_log()
    info_log = linker.get_info_log()
    assert isinstance(err_log, str)
    assert isinstance(info_log, str)
    # Calling again should return the same observable values.
    assert linker.get_error_log() == err_log
    assert linker.get_info_log() == info_log


def test_linker_handle(compile_ptx_functions):
    """Linker.handle returns a non-null handle object."""
    options = LinkerOptions(arch=ARCH)
    linker = Linker(*compile_ptx_functions, options=options)
    handle = linker.handle
    assert handle is not None
    assert int(handle) != 0


@pytest.mark.agent_authored(model="gpt-5")
@pytest.mark.skipif(not is_culink_backend, reason="driver backend regression test")
def test_driver_linker_lifetime_no_heap_corruption(compile_ptx_functions):
    if not _can_load_generated_ptx():
        pytest.skip("PTX version too new for current driver")

    linker = Linker(*compile_ptx_functions, options=LinkerOptions(arch=ARCH))
    linker.link("cubin")
    linker.close()
    del linker

    obj_a = Program(kernel_a, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    obj_b = Program(device_function_b, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    obj_c = Program(device_function_c, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    linker = Linker(obj_a, obj_b, obj_c, options=LinkerOptions(arch=ARCH))
    linker.link("cubin")
    linker.close()


@pytest.mark.agent_authored(model="gpt-5")
@pytest.mark.skipif(not is_culink_backend, reason="driver backend regression test")
def test_driver_linker_preserves_error_log_after_close(init_cuda):
    if not _can_load_generated_ptx():
        pytest.skip("PTX version too new for current driver")

    bad_kernel = """
extern __device__ int Z();
__global__ void A() { int r = Z(); }
"""
    bad_obj = Program(bad_kernel, "c++", ProgramOptions(relocatable_device_code=True)).compile("ptx")
    linker = Linker(bad_obj, options=LinkerOptions(arch=ARCH))
    with pytest.raises(CUDAError):
        linker.link("cubin")

    error_log = linker.get_error_log()
    assert error_log
    linker.close()
    assert linker.get_error_log() == error_log
    assert isinstance(linker.get_info_log(), str)


@pytest.mark.skipif(is_culink_backend, reason="nvjitlink options only tested with nvjitlink backend")
def test_linker_options_nvjitlink_options_as_str():
    """_prepare_nvjitlink_options(as_bytes=False) returns plain strings."""
    opts = LinkerOptions(arch=ARCH, debug=True, lineinfo=True)
    options = opts._prepare_nvjitlink_options(as_bytes=False)
    assert isinstance(options, list)
    assert all(isinstance(o, str) for o in options)
    assert f"-arch={ARCH}" in options
    assert "-g" in options
    assert "-lineinfo" in options


class TestWhichBackendClassmethod:
    def test_which_backend_returns_nvjitlink(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", True)
        assert Linker.which_backend() == "nvJitLink"

    def test_which_backend_returns_driver(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", False)
        assert Linker.which_backend() == "driver"

    def test_which_backend_invokes_probe_when_not_memoised(self, monkeypatch):
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", None)
        called = []

        def fake_decide():
            called.append(True)
            return False  # False = not falling back to driver = nvJitLink

        monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", fake_decide)
        result = Linker.which_backend()
        assert result == "nvJitLink"
        assert called, "_decide_nvjitlink_or_driver was not called"

    @pytest.mark.agent_authored(model="grok-4.5")
    def test_which_backend_falls_back_when_nvjitlink_too_old(self, monkeypatch):
        """Regression test for #2408: old nvJitLink must not crash which_backend()."""
        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", None)
        monkeypatch.setattr(_linker, "_driver", None)

        def fake__optional_cuda_import(modname, probe_function=None):
            assert modname == "cuda.bindings.nvjitlink"
            assert probe_function is None
            return object()

        monkeypatch.setattr(_linker, "_optional_cuda_import", fake__optional_cuda_import)
        monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", lambda _nvjitlink: False)

        with pytest.warns(RuntimeWarning, match="too old \\(<12.3\\)"):
            assert Linker.which_backend() == "driver"

        assert _linker._use_nvjitlink_backend is False

    @pytest.mark.agent_authored(model="grok-4.5")
    def test_which_backend_falls_back_when_dylib_missing(self, monkeypatch):
        """Missing nvJitLink dylib must fall back without raising."""
        from cuda.pathfinder import DynamicLibNotFoundError

        monkeypatch.setattr(_linker, "_use_nvjitlink_backend", None)
        monkeypatch.setattr(_linker, "_driver", None)

        def raise_missing(_nvjitlink):
            raise DynamicLibNotFoundError("missing")

        def fake__optional_cuda_import(modname, probe_function=None):
            assert modname == "cuda.bindings.nvjitlink"
            assert probe_function is None
            return object()

        monkeypatch.setattr(_linker, "_nvjitlink_has_version_symbol", raise_missing)
        monkeypatch.setattr(_linker, "_optional_cuda_import", fake__optional_cuda_import)

        with pytest.warns(RuntimeWarning, match="cuda.bindings.nvjitlink is not available"):
            assert Linker.which_backend() == "driver"

        assert _linker._use_nvjitlink_backend is False

    def test_which_backend_is_classmethod(self):
        attr = inspect.getattr_static(Linker, "which_backend")
        assert isinstance(attr, classmethod)

    def test_which_backend_is_not_property(self):
        """which_backend is a classmethod, not a property.

        This is an intentional breaking change from the prior ``backend`` property API.
        All call sites must use parens: ``Linker.which_backend()``.
        """
        attr = inspect.getattr_static(Linker, "which_backend")
        assert not isinstance(attr, property)


@pytest.fixture
def driver_binding(monkeypatch):
    """Pin _linker._driver to the real driver module so driver-backend tests run under any backend."""
    from cuda.bindings import driver

    monkeypatch.setattr(_linker, "_driver", driver)
    return driver


def test_prepare_driver_options_all_supported(driver_binding):
    """Exercise every supported branch of _prepare_driver_options."""
    driver = driver_binding
    opts = LinkerOptions(
        arch="sm_80",
        max_register_count=32,
        verbose=True,
        link_time_optimization=True,
        optimization_level=2,
        debug=True,
        lineinfo=True,
        no_cache=True,
    )
    formatted, keys = opts._prepare_driver_options()
    assert len(formatted) == len(keys)
    assert len(keys) == 4 + 8  # 4 fixed log-buffer entries + 8 options set above

    # Skip log-buffer entries; verify key-to-value mapping (catches swap/dup/wrong-value).
    payload_keys = keys[4:]
    assert len(set(payload_keys)) == len(payload_keys), f"duplicate option keys: {payload_keys}"
    option_to_value = dict(zip(payload_keys, formatted[4:]))
    assert option_to_value[driver.CUjit_option.CU_JIT_TARGET] == driver.CUjit_target.CU_TARGET_COMPUTE_80
    assert option_to_value[driver.CUjit_option.CU_JIT_MAX_REGISTERS] == 32
    assert option_to_value[driver.CUjit_option.CU_JIT_LOG_VERBOSE] == 1
    assert option_to_value[driver.CUjit_option.CU_JIT_LTO] == 1
    assert option_to_value[driver.CUjit_option.CU_JIT_OPTIMIZATION_LEVEL] == 2
    assert option_to_value[driver.CUjit_option.CU_JIT_GENERATE_DEBUG_INFO] == 1
    assert option_to_value[driver.CUjit_option.CU_JIT_GENERATE_LINE_INFO] == 1
    assert option_to_value[driver.CUjit_option.CU_JIT_CACHE_MODE] == driver.CUjit_cacheMode.CU_JIT_CACHE_OPTION_NONE


@pytest.mark.parametrize(
    "kwargs,match",
    [
        ({"ftz": True}, "ftz option is deprecated"),
        ({"prec_div": True}, "prec_div option is deprecated"),
        ({"prec_sqrt": True}, "prec_sqrt option is deprecated"),
        ({"fma": True}, "fma options is deprecated"),
        ({"kernels_used": "my_kernel"}, "kernels_used is deprecated"),
        ({"variables_used": "my_var"}, "variables_used is deprecated"),
        ({"optimize_unused_variables": True}, "optimize_unused_variables is deprecated"),
    ],
)
def test_prepare_driver_options_deprecated_warnings(driver_binding, kwargs, match):
    """Each driver-deprecated option emits a DeprecationWarning."""
    opts = LinkerOptions(**kwargs)
    with pytest.warns(DeprecationWarning, match=match):
        opts._prepare_driver_options()


@pytest.mark.parametrize(
    "kwargs,match",
    [
        ({"time": True}, "time option is not supported by the driver API"),
        ({"ptx": True}, "ptx option is not supported by the driver API"),
        ({"ptxas_options": ["-v"]}, "ptxas_options option is not supported by the driver API"),
        ({"split_compile": 0}, "split_compile option is not supported by the driver API"),
        ({"split_compile_extended": 1}, "split_compile_extended option is not supported by the driver API"),
    ],
)
def test_prepare_driver_options_unsupported_raises(driver_binding, kwargs, match):
    """Each nvjitlink-only option raises ValueError on the driver backend."""
    opts = LinkerOptions(**kwargs)
    with pytest.raises(ValueError, match=match):
        opts._prepare_driver_options()


def test_linker_empty_object_codes_raises():
    """Linker with no ObjectCode raises ValueError."""
    with pytest.raises(ValueError, match="At least one ObjectCode object must be provided"):
        Linker()


def test_as_bytes_nvjitlink_unavailable(monkeypatch):
    """as_bytes('nvjitlink') raises RuntimeError when the backend is unavailable."""
    monkeypatch.setattr(_linker, "_use_nvjitlink_backend", False)
    opts = LinkerOptions(arch="sm_80")
    with pytest.raises(RuntimeError, match="nvJitLink backend is not available"):
        opts.as_bytes("nvjitlink")
