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


import base64
import binascii
import shutil
import subprocess

import pytest

from cuda.bindings import nvfatbin
from cuda.bindings._v2 import nvrtc

ARCHITECTURES = ["sm_75", "sm_80", "sm_90", "sm_100"]
PTX_VERSIONS = ["6.4", "7.0", "8.5", "8.8"]

PTX_TEMPLATE = """
.version {PTX_VERSION}
.target {ARCH}
.address_size 64

        // .globl       _Z6kernelPi

.visible .entry _Z6kernelPi(
        .param .u64 _Z6kernelPi_param_0
)
{{
        .reg .b32       %r<7>;
        .reg .b64       %rd<5>;


        ld.param.u64    %rd1, [_Z6kernelPi_param_0];
        cvta.to.global.u64      %rd2, %rd1;
        mov.u32         %r1, %tid.x;
        mov.u32         %r2, %ctaid.x;
        mov.u32         %r3, %ntid.x;
        mad.lo.s32      %r4, %r2, %r3, %r1;
        mul.wide.s32    %rd3, %r4, 4;
        add.s64         %rd4, %rd2, %rd3;
        ld.global.u32   %r5, [%rd4];
        add.s32         %r6, %r5, 1;
        st.global.u32   [%rd4], %r6;
        ret;

}}
"""

CODE = """
int __device__ inc(int x) {
    return x + 1;
}
"""

# Base64 encoded TileIR generated by the toolshed/dump_cutile_b64.py script.
TILEIR_b64 = (
    "f1RpbGVJUgANAQAAgo0BCAECBgYBCwEECgCBAUQHBgQIEAAABgUMAQABBgUIEAALQwEICgEMAAYE"
    "CBAAAwYFDAEABAYFCBAAD0MBCA4BEAAGBAgQAAYGBQwBAAcGBQgQABNDAQgSARQAMAUFBUIJDT4C"
    "CgcEABkBFglCCRE+AgoHBAAcARYJAgoAABodQgkVZgEHBAAfIAEWCVwAAIQICADLy8vLy8vLg5oC"
    "CMvLy8sBy8vLAAAAABfLy8vLy8vLBAAAAAAAAAAEAAAAAAAAAAQAAAAAAAAABAAAAAAAAAAEAAAA"
    "AAAAAAQAAAAAAAAABAAAAAAAAAAEAAAAAAAAAAQAAAAAAAAABAAAAAAAAAAEAAAAAAAAAAQAAAAA"
    "AAAABAAAAAAAAAAEAAAAAAAAAAUAAAAAAAAABgAAAAAAAAAGAAAAAAAAAAcAAAAAAAAABwAAAAAA"
    "AAAIAAAAAAAAAAkAAAAAAAAACQAAAAAAAAAEAAAAAAAAAAnLy8sAAAAAAwAAAAUAAAAMAAAAEQAA"
    "ABYAAAAbAAAAIAAAACUAAAACAAEBAQUBFwICAhcEAwMYBAQDAxkTBAMDGhEEAwMbEQQDAx0WBAMD"
    "HgiFdATLy8sLy8vLAAAAAAEAAAACAAAAAwAAAAUAAAAIAAAACwAAABcAAAAYAAAALAAAADkAAAAA"
    "AwcMAg0DAA0BABAJBAUFBAUFBAUFABEOAgEAAAAAAAAAgAEBAAAAAAAAAA8BEAAAAAgBAAAAAAAN"
    "AgEQAAAAAAAAAIGIAQQFy8vLAAAAABIAAAAsAAAAPQAAAGoAAABkdW1wX2N1dGlsZV9iNjQucHkv"
    "bG9jYWxob21lL2xvY2FsLXdhbmdtL3RveXZlY3Rvcl9hZGRfa2VybmVsL2xvY2FsaG9tZS9sb2Nh"
    "bC13YW5nbS90b3kvZHVtcF9jdXRpbGVfYjY0LnB5c21fMTIwAA=="
)


def get_version() -> tuple[int, int]:
    return nvfatbin.version()


@pytest.fixture(params=ARCHITECTURES)
def arch(request):
    return request.param


@pytest.fixture(params=PTX_VERSIONS)
def ptx_version(request):
    return request.param


@pytest.fixture
def PTX(arch, ptx_version):
    return PTX_TEMPLATE.format(PTX_VERSION=ptx_version, ARCH=arch)


@pytest.fixture
def nvcc_smoke(tmpdir) -> str:
    # TODO: Use cuda-pathfinder to locate nvcc on system.
    nvcc = shutil.which("nvcc")
    if nvcc is None:
        pytest.skip("nvcc not found on PATH")

    # Smoke test: make sure nvcc is actually usable (toolkit + host compiler are set up),
    # not merely present on PATH.
    src = tmpdir / "nvcc_smoke.cu"
    out = tmpdir / "nvcc_smoke.o"
    with open(src, "w") as f:
        f.write("")
    try:
        subprocess.run(  # noqa: S603
            [nvcc, "-c", str(src), "-o", str(out)],
            check=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as e:
        stdout = (e.stdout or b"").decode(errors="replace")
        stderr = (e.stderr or b"").decode(errors="replace")
        pytest.skip(
            "nvcc found on PATH but failed to compile a trivial input.\n"
            f"command: {[nvcc, '-c', str(src), '-o', str(out)]!r}\n"
            f"exit_code: {e.returncode}\n"
            f"stdout:\n{stdout}\n"
            f"stderr:\n{stderr}\n"
        )

    return nvcc


def _build_cubin(arch):
    program_handle = nvrtc.create_program(CODE.encode(), b"")
    nvrtc.compile_program(program_handle, [f"-arch={arch}".encode()])
    return nvrtc.get_cubin(program_handle)


@pytest.fixture
def CUBIN(arch):
    return _build_cubin(arch)


# create a valid LTOIR input for testing
@pytest.fixture
def LTOIR(arch):
    arch = arch.replace("sm", "compute")
    empty_cplusplus_kernel = "__global__ void A() {}"
    program_handle = nvrtc.create_program(empty_cplusplus_kernel.encode(), b"")
    nvrtc.compile_program(program_handle, [b"-dlto", f"-arch={arch}".encode()])
    return nvrtc.get_ltoir(program_handle)


@pytest.fixture
def OBJECT(arch, tmpdir, nvcc_smoke):
    empty_cplusplus_kernel = "__global__ void A() {}"
    with open(tmpdir / "object.cu", "w") as f:
        f.write(empty_cplusplus_kernel)

    nvcc = nvcc_smoke

    # This is a test fixture that intentionally invokes a trusted tool (`nvcc`) to
    # compile a temporary CUDA translation unit.
    cmd = [nvcc, "-c", "-arch", arch, "-o", str(tmpdir / "object.o"), str(tmpdir / "object.cu")]
    try:
        subprocess.run(  # noqa: S603
            cmd,
            check=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError as e:
        stdout = (e.stdout or b"").decode(errors="replace")
        stderr = (e.stderr or b"").decode(errors="replace")
        raise RuntimeError(
            "nvcc smoke test passed, but nvcc failed while compiling the test object.\n"
            f"command: {cmd!r}\n"
            f"exit_code: {e.returncode}\n"
            f"stdout:\n{stdout}\n"
            f"stderr:\n{stderr}\n"
        ) from e
    with open(tmpdir / "object.o", "rb") as f:
        object = f.read()

    return object


@pytest.fixture
def TILEIR(tmpdir):
    try:
        binary_data = base64.b64decode(TILEIR_b64)
    except binascii.Error as e:
        raise ValueError(
            "Base64 encoded TileIR is corrupted. Please regenerate the TileIR"
            "by executing the toolshed/dump_cutile_b64.py script."
        ) from e
    return binary_data


@pytest.mark.parametrize("error_enum", nvfatbin.Result)
def test_get_error_string(error_enum):
    es = nvfatbin.get_error_string(error_enum)

    if error_enum is nvfatbin.Result.SUCCESS:
        assert es == ""
    else:
        assert es != ""


def test_nvfatbin_get_version():
    major, minor = nvfatbin.version()
    assert major is not None
    assert minor is not None


def test_nvfatbin_empty_create_and_destroy():
    handle = nvfatbin.create([], 0)
    assert handle is not None
    nvfatbin.destroy(handle)


def test_nvfatbin_invalid_input_create():
    with pytest.raises(nvfatbin.nvFatbinError, match="ERROR_UNRECOGNIZED_OPTION"):
        nvfatbin.create(["--unsupported_option"], 1)


def test_nvfatbin_get_empty():
    handle = nvfatbin.create([], 0)
    size = nvfatbin.size(handle)

    buffer = bytearray(size)
    nvfatbin.get(handle, buffer)

    nvfatbin.destroy(handle)


def test_nvfatbin_add_ptx(PTX, arch):
    arch_numeric = arch.split("_")[1]

    handle = nvfatbin.create([], 0)
    nvfatbin.add_ptx(handle, PTX.encode(), len(PTX), arch_numeric, "add", f"-arch={arch}")

    buffer = bytearray(nvfatbin.size(handle))

    nvfatbin.get(handle, buffer)
    nvfatbin.destroy(handle)


def test_nvfatbin_add_cubin_ELF_SIZE_MISMATCH():
    cubin = _build_cubin("sm_80")
    handle = nvfatbin.create([], 0)
    with pytest.raises(nvfatbin.nvFatbinError, match="ERROR_ELF_ARCH_MISMATCH"):
        nvfatbin.add_cubin(handle, cubin, len(cubin), "75", "inc")

    nvfatbin.destroy(handle)


def test_nvfatbin_add_cubin(CUBIN, arch):
    arch_numeric = arch.split("_")[1]

    handle = nvfatbin.create([], 0)
    nvfatbin.add_cubin(handle, CUBIN, len(CUBIN), arch_numeric, "inc")

    buffer = bytearray(nvfatbin.size(handle))

    nvfatbin.get(handle, buffer)
    nvfatbin.destroy(handle)


def test_nvfatbin_add_cubin_ELF_ARCH_MISMATCH():
    cubin = _build_cubin("sm_80")
    handle = nvfatbin.create([], 0)
    with pytest.raises(nvfatbin.nvFatbinError, match="ERROR_ELF_ARCH_MISMATCH"):
        nvfatbin.add_cubin(handle, cubin, len(cubin), "75", "inc")

    nvfatbin.destroy(handle)


def test_nvdfatbin_add_ltoir(LTOIR, arch):
    arch_numeric = arch.split("_")[1]

    handle = nvfatbin.create([], 0)
    nvfatbin.add_ltoir(handle, LTOIR, len(LTOIR), arch_numeric, "inc", "")

    buffer = bytearray(nvfatbin.size(handle))

    nvfatbin.get(handle, buffer)
    nvfatbin.destroy(handle)


def test_nvfatbin_add_reloc(OBJECT):
    handle = nvfatbin.create([], 0)
    nvfatbin.add_reloc(handle, OBJECT, len(OBJECT))

    buffer = bytearray(nvfatbin.size(handle))

    nvfatbin.get(handle, buffer)
    nvfatbin.destroy(handle)


@pytest.mark.skipif(get_version() < (13, 1), reason="TileIR API is not supported in CUDA < 13.1")
def test_nvfatbin_add_tile_ir(TILEIR):
    handle = nvfatbin.create([], 0)
    nvfatbin.add_tile_ir(handle, TILEIR, len(TILEIR), "VectorAdd", "")

    buffer = bytearray(nvfatbin.size(handle))

    nvfatbin.get(handle, buffer)
    nvfatbin.destroy(handle)
