# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Build and load C++/CUDA sources into a tvm_ffi Module using Ninja."""

from __future__ import annotations

import functools
import hashlib
import os
import shutil
import subprocess
import sys
from collections.abc import Mapping, Sequence
from contextlib import nullcontext
from pathlib import Path
from typing import Any, Literal

from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path, find_libtvm_ffi
from tvm_ffi.module import Module, load_module
from tvm_ffi.utils import FileLock

IS_WINDOWS = sys.platform == "win32"
BACKEND_STR = Literal["cuda", "hip"]


@functools.lru_cache
def _detect_gpu_backend() -> BACKEND_STR:
    """Auto-detect whether to use CUDA or HIP (ROCm).

    Returns 'hip' if ROCm/HIP is available, 'cuda' otherwise.
    """
    # Check environment variable override first
    backend = os.environ.get("TVM_FFI_GPU_BACKEND", "").lower()
    if backend in ("cuda", "hip"):
        return backend  # type: ignore[return-value]
    try:
        _find_rocm_home()
        return "hip"
    except RuntimeError:
        return "cuda"


def _resolve_gpu_backend(backend: str | None) -> BACKEND_STR:
    if backend is not None:
        if backend in ("cuda", "hip"):
            return backend  # type: ignore[return-value]
        raise ValueError(f"Invalid backend: {backend}. Supported backends are 'cuda' and 'hip'.")
    return _detect_gpu_backend()


def _hash_sources(
    cpp_source: str | None,
    cuda_source: str | None,
    cpp_files: Sequence[str] | None,
    cuda_files: Sequence[str] | None,
    functions: Sequence[str] | Mapping[str, str],
    extra_cflags: Sequence[str],
    extra_cuda_cflags: Sequence[str],
    extra_ldflags: Sequence[str],
    extra_include_paths: Sequence[str],
    embed_cubin: Mapping[str, bytes] | None = None,
) -> str:
    """Generate a unique hash for the given sources and functions."""
    m = hashlib.sha256()

    def _hash(obj: Any) -> None:
        if obj is None:
            m.update(b"None")
        elif isinstance(obj, str):
            m.update(b"str")
            m.update(obj.encode("utf-8"))
        elif isinstance(obj, bytes):
            m.update(b"bytes")
            m.update(obj)
        elif isinstance(obj, Mapping):
            m.update(b"Mapping")
            for key in sorted(obj.keys()):
                item = obj[key]
                _hash(key)
                _hash(item)
        elif isinstance(obj, Sequence):
            m.update(b"Sequence")
            for item in obj:
                _hash(item)
        else:
            raise ValueError(f"Unsupported type: {type(obj)}")

    _hash(
        (
            cpp_source,
            cuda_source,
            sorted(cpp_files) if cpp_files is not None else None,
            sorted(cuda_files) if cuda_files is not None else None,
            functions,
            extra_cflags,
            extra_cuda_cflags,
            extra_ldflags,
            extra_include_paths,
            embed_cubin,
        )
    )

    return m.hexdigest()[:16]


def _maybe_write(path: str, content: str) -> None:
    """Write content to path if it does not already exist with the same content."""
    p = Path(path)
    if p.exists():
        with p.open() as f:
            existing_content = f.read()
        if existing_content == content:
            return
    with p.open("w") as f:
        f.write(content)


@functools.lru_cache
def _find_cuda_home() -> str:
    """Find the CUDA install path."""
    # Guess #1
    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
    if cuda_home is None:
        # Guess #2
        nvcc_path = shutil.which("nvcc")
        if nvcc_path is not None:
            cuda_home = str(Path(nvcc_path).parent.parent)
        else:
            # Guess #3
            if IS_WINDOWS:
                cuda_root = Path("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA")
                cuda_homes = list(cuda_root.glob("v*.*"))
                if len(cuda_homes) == 0:
                    raise RuntimeError(
                        "Could not find CUDA installation. Please set CUDA_HOME environment variable."
                    )
                cuda_home = str(cuda_homes[0])
            else:
                cuda_home = "/usr/local/cuda"
            if not Path(cuda_home).exists():
                raise RuntimeError(
                    "Could not find CUDA installation. Please set CUDA_HOME environment variable."
                )
    return cuda_home


def _get_cuda_target() -> str:
    """Get the CUDA target architecture flag."""
    if "TVM_FFI_CUDA_ARCH_LIST" in os.environ:
        arch_list = os.environ["TVM_FFI_CUDA_ARCH_LIST"].split()  # e.g., "8.9 9.0a"
        flags = []
        for arch in arch_list:
            if len(arch.split(".")) != 2:
                raise ValueError(f"Invalid CUDA architecture: {arch}")
            major, minor = arch.split(".")
            flags.append(f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}")
        return " ".join(flags)
    else:
        try:
            status = subprocess.run(
                args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
                capture_output=True,
                check=True,
            )
            compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
            major, minor = compute_cap.split(".")
            return f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"
        except Exception:
            try:
                # For old drivers, there is no compute_cap, but we can use the GPU name to determine the architecture.
                ampere_arch_map = {
                    "A100": ("8", "0"),
                    "A10": ("8", "6"),
                }
                status = subprocess.run(
                    args=["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
                    capture_output=True,
                    check=True,
                    text=True,
                )
                gpu_name = status.stdout.strip().split("\n")[0]
                for gpu_key, (major, minor) in ampere_arch_map.items():
                    if gpu_key in gpu_name:
                        return f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"
            except (subprocess.CalledProcessError, FileNotFoundError):
                pass
            raise RuntimeError(
                "Could not detect CUDA compute_cap automatically. Please set TVM_FFI_CUDA_ARCH_LIST environment variable."
            )


@functools.lru_cache
def _find_rocm_home() -> str:
    """Find the ROCm install path."""
    # Guess #1: check environment variables
    rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH")
    if rocm_home is None:
        hipcc_path = shutil.which("hipcc")
        # Guess #2: find hipcc in PATH and resolve ROCm home from it
        if hipcc_path is not None:
            rocm_home = str(Path(hipcc_path).resolve().parent.parent)
            if Path(rocm_home).name == "hip":
                rocm_home = str(Path(rocm_home).parent)
        else:
            # Guess #3: use default installation path
            rocm_home = "/opt/rocm"
            if not Path(rocm_home).exists():
                raise RuntimeError(
                    "Could not find ROCm installation. Please set ROCM_HOME environment variable."
                )
    return rocm_home


def _get_rocm_target() -> list[str]:
    """Get the ROCm target architecture flags (--offload-arch=gfxXXXX)."""
    if "TVM_FFI_ROCM_ARCH_LIST" in os.environ:
        arch_list = os.environ["TVM_FFI_ROCM_ARCH_LIST"].split()  # e.g., "gfx90a gfx942"
        return [f"--offload-arch={arch}" for arch in arch_list]
    # Try rocm_agent_enumerator
    try:
        agent_enum = str(Path(_find_rocm_home()) / "bin" / "rocm_agent_enumerator")
        if not Path(agent_enum).exists():
            agent_enum = "rocm_agent_enumerator"
        status = subprocess.run(args=[agent_enum], capture_output=True, check=True, text=True)
        archs = list(
            dict.fromkeys(
                line.strip()
                for line in status.stdout.strip().split("\n")
                if line.strip() and line.strip() != "gfx000"
            )
        )
        if archs:
            return [f"--offload-arch={arch}" for arch in archs]
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass
    # Try rocminfo
    try:
        status = subprocess.run(args=["rocminfo"], capture_output=True, check=True, text=True)
        archs = list(
            dict.fromkeys(
                line.split(":")[-1].strip()
                for line in status.stdout.split("\n")
                if "Name:" in line
                and "gfx" in line.lower()
                and line.split(":")[-1].strip() != "gfx000"
            )
        )
        if archs:
            return [f"--offload-arch={arch}" for arch in archs]
    except (subprocess.CalledProcessError, FileNotFoundError):
        pass
    raise RuntimeError(
        "Could not detect ROCm GPU architecture automatically. "
        "Please set TVM_FFI_ROCM_ARCH_LIST environment variable (e.g. 'gfx90a gfx942')."
    )


def _run_command_in_dev_prompt(
    args: list[str],
    cwd: str | os.PathLike[str],
    capture_output: bool,
) -> subprocess.CompletedProcess:
    """Locates the Developer Command Prompt and runs a command within its environment."""
    try:
        # Path to vswhere.exe
        vswhere_path = str(
            Path(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)"))
            / "Microsoft Visual Studio"
            / "Installer"
            / "vswhere.exe"
        )

        if not Path(vswhere_path).exists():
            raise FileNotFoundError("vswhere.exe not found.")

        # Find the Visual Studio installation path
        vs_install_path = subprocess.run(
            [
                vswhere_path,
                "-latest",
                "-prerelease",
                "-products",
                "*",
                "-property",
                "installationPath",
            ],
            capture_output=True,
            text=True,
            check=True,
        ).stdout.strip()

        if not vs_install_path:
            raise FileNotFoundError("No Visual Studio installation found.")

        # Construct the path to the VsDevCmd.bat file
        vsdevcmd_path = str(Path(vs_install_path) / "Common7" / "Tools" / "VsDevCmd.bat")

        if not Path(vsdevcmd_path).exists():
            raise FileNotFoundError(f"VsDevCmd.bat not found at: {vsdevcmd_path}")

        # Use cmd.exe to run the batch file and then your command.
        # The /k flag keeps the command prompt open after the batch file runs.
        # The "&" symbol chains the commands.
        cmd_command = '"{vsdevcmd_path}" -arch=x64 & {command}'.format(
            vsdevcmd_path=vsdevcmd_path, command=" ".join(args)
        )

        # Execute the command in a new shell
        return subprocess.run(
            cmd_command, check=False, cwd=cwd, capture_output=capture_output, shell=True
        )

    except (FileNotFoundError, subprocess.CalledProcessError) as e:
        raise RuntimeError(
            "Failed to run the following command in MSVC developer environment: {}".format(
                " ".join(args)
            )
        ) from e


def _generate_ninja_build(  # noqa: PLR0915, PLR0912
    name: str,
    extra_cflags: Sequence[str],
    extra_cuda_cflags: Sequence[str],
    extra_ldflags: Sequence[str],
    extra_include_paths: Sequence[str],
    sources: Sequence[str],
    embed_cubin: Mapping[str, bytes] | None = None,
    backend: str | None = None,
    output: str | None = None,
) -> str:
    """Generate the content of build.ninja for building the module."""
    # Determine output format from extension
    if output is not None:
        out_ext = Path(output).suffix.lower()
        object_mode = out_ext in (".o", ".obj")
        output_name = output
    else:
        object_mode = False
        output_name = f"{name}{'.dll' if IS_WINDOWS else '.so'}"
    has_cuda_sources = any(Path(s).suffix.lower() == ".cu" for s in sources)
    with_hip = backend == "hip"
    with_cuda = backend == "cuda"
    with_backend = with_hip or with_cuda or has_cuda_sources
    if has_cuda_sources and not (with_hip or with_cuda):
        # Auto-detect backend from available GPU
        detected = _resolve_gpu_backend(None)
        with_hip = detected == "hip"
        with_cuda = detected == "cuda"

    default_include_paths = [find_include_path(), find_dlpack_include_path()]
    tvm_ffi_lib = Path(find_libtvm_ffi())
    tvm_ffi_lib_path = str(tvm_ffi_lib.parent)
    tvm_ffi_lib_name = tvm_ffi_lib.stem
    if IS_WINDOWS:
        default_cflags = ["/O2", "/MD"]
        default_cxxflags = ["/std:c++17", "/MD", "/EHsc"]
        _win_warnings = [
            "/wd4819",
            "/wd4251",
            "/wd4244",
            "/wd4267",
            "/wd4275",
            "/wd4018",
            "/wd4190",
            "/wd4624",
            "/wd4067",
            "/wd4068",
        ]
        default_cflags += _win_warnings
        default_cxxflags += _win_warnings
        default_cuda_cflags = ["-Xcompiler", "/std:c++17", "/O2"]
        default_ldflags = [
            "/DLL",
            f"/LIBPATH:{tvm_ffi_lib_path}",
            f"{tvm_ffi_lib_name}.lib",
        ]
    else:
        default_cflags = ["-fPIC", "-O2"]
        default_cxxflags = ["-std=c++17", "-fPIC", "-O2"]
        default_cuda_cflags = ["-std=c++17", "-O2"]
        default_ldflags = ["-shared", f"-L{tvm_ffi_lib_path}", "-ltvm_ffi"]

        if with_hip:
            rocm_home = _find_rocm_home()
            default_cuda_cflags += ["-fPIC", "-D__HIP_PLATFORM_AMD__=1", "-fno-gpu-rdc"]
            default_cuda_cflags += _get_rocm_target()
            default_include_paths.append(str(Path(rocm_home) / "include"))
            default_ldflags += [
                f"-L{Path(rocm_home) / 'lib'!s}",
                "-lamdhip64",
            ]
        if with_cuda:
            default_cuda_cflags = ["-Xcompiler", "-fPIC", *default_cuda_cflags]
            default_cuda_cflags += [_get_cuda_target()]
            default_ldflags += [
                "-L{}".format(str(Path(_find_cuda_home()) / "lib64")),
                "-lcudart",  # cuda runtime library
            ]

    extra_cflags_list = [flag.strip() for flag in extra_cflags]
    cflags = default_cflags + extra_cflags_list
    cxxflags = default_cxxflags + extra_cflags_list
    cuda_cflags = default_cuda_cflags + [flag.strip() for flag in extra_cuda_cflags]
    ldflags = default_ldflags + [flag.strip() for flag in extra_ldflags]
    include_paths = default_include_paths + [
        str(Path(path).resolve()) for path in extra_include_paths
    ]

    # append include paths
    for path in include_paths:
        inc = "-I{}".format(path.replace(":", "$:"))
        cflags.append(inc)
        cxxflags.append(inc)
        cuda_cflags.append(inc)

    # Classify sources by extension to determine which rules are needed
    with_c = any(Path(s).suffix.lower() == ".c" for s in sources)
    ninja: list[str] = []
    ninja.append("ninja_required_version = 1.3")
    ninja.append("cxx = {}".format(os.environ.get("CXX", "cl" if IS_WINDOWS else "c++")))
    ninja.append("cxxflags = {}".format(" ".join(cxxflags)))
    if with_c:
        ninja.append("cc = {}".format(os.environ.get("CC", "cl" if IS_WINDOWS else "cc")))
        ninja.append("cflags = {}".format(" ".join(cflags)))
    if with_backend:
        if with_hip:
            ninja.append("nvcc = {}".format(str(Path(_find_rocm_home()) / "bin" / "hipcc")))
        if with_cuda:
            ninja.append("nvcc = {}".format(str(Path(_find_cuda_home()) / "bin" / "nvcc")))
        ninja.append("cuda_cflags = {}".format(" ".join(cuda_cflags)))
    ninja.append("ldflags = {}".format(" ".join(ldflags)))

    # rules
    ninja.append("")
    ninja.append("rule compile")
    if IS_WINDOWS:
        ninja.append("  command = $cxx /showIncludes $cxxflags -c $in /Fo$out")
        ninja.append("  deps = msvc")
    else:
        ninja.append("  depfile = $out.d")
        ninja.append("  deps = gcc")
        ninja.append("  command = $cxx -MMD -MF $out.d $cxxflags -c $in -o $out")
    ninja.append("")

    if with_c:
        ninja.append("rule c_compile")
        if IS_WINDOWS:
            ninja.append("  command = $cc /showIncludes $cflags -c $in /Fo$out")
            ninja.append("  deps = msvc")
        else:
            ninja.append("  depfile = $out.d")
            ninja.append("  deps = gcc")
            ninja.append("  command = $cc -MMD -MF $out.d $cflags -c $in -o $out")
        ninja.append("")

    if with_backend:
        ninja.append("rule compile_cuda")
        ninja.append("  depfile = $out.d")
        ninja.append("  deps = gcc")
        if with_hip:
            ninja.append("  command = $nvcc $cuda_cflags -c $in -o $out")
        else:
            ninja.append(
                "  command = $nvcc  --generate-dependencies-with-compile --dependency-output $out.d $cuda_cflags -c $in -o $out"
            )
        ninja.append("")

    # Add rules for object merging and cubin embedding (Unix only)
    if not IS_WINDOWS:
        ninja.append("rule merge_objects")
        ninja.append("  command = ld -r -o $out $in")
        ninja.append("")

        if embed_cubin:
            ninja.append("rule embed_cubin")
            ninja.append(
                f"  command = {sys.executable} -m tvm_ffi.utils.embed_cubin --output-obj $out --input-obj $in --cubin $cubin --name $name"
            )
            ninja.append("")

    if not object_mode:
        ninja.append("rule link")
        if IS_WINDOWS:
            ninja.append("  command = $cxx $in /link $ldflags /out:$out")
        else:
            ninja.append("  command = $cxx $in $ldflags -o $out")
        ninja.append("")

    # build targets — dispatch by file extension
    obj_files: list[str] = []
    c_idx = cpp_idx = cuda_idx = 0
    for src in sorted(sources):
        ext = Path(src).suffix.lower()
        escaped = src.replace(":", "$:")
        if ext in (".o", ".obj"):
            # Pre-compiled object file: pass directly to linker
            obj_files.append(escaped)
        elif ext == ".c":
            obj_name = f"c_{c_idx}.o"
            ninja.append(f"build {obj_name}: c_compile {escaped}")
            obj_files.append(obj_name)
            c_idx += 1
        elif ext == ".cu":
            obj_name = f"cuda_{cuda_idx}.o"
            ninja.append(f"build {obj_name}: compile_cuda {escaped}")
            obj_files.append(obj_name)
            cuda_idx += 1
        else:
            # .cc, .cpp, .cxx — default to C++ compilation
            obj_name = f"cpp_{cpp_idx}.o"
            ninja.append(f"build {obj_name}: compile {escaped}")
            obj_files.append(obj_name)
            cpp_idx += 1

    if object_mode:
        # Object-only output: merge all object files into the target.
        if not IS_WINDOWS:
            ninja.append(f"build {output_name}: merge_objects {' '.join(obj_files)}")
            ninja.append("")
            ninja.append(f"default {output_name}")
        else:
            # Windows: no ld -r available; default to the first intermediate object
            ninja.append(f"default {obj_files[0]}")
        ninja.append("")
        return "\n".join(ninja)

    # For Unix systems with embed_cubin, use a 3-step process:
    # 1. Merge all object files into a unified object file
    # 2. Embed each cubin into the unified object file (chain them)
    # 3. Link the final object file into a shared library
    if not IS_WINDOWS and embed_cubin:
        # Step 1: Merge object files into unified.o
        unified_obj = "unified.o"
        obj_files_str = " ".join(obj_files)
        ninja.append(f"build {unified_obj}: merge_objects {obj_files_str}")
        ninja.append("")

        # Step 2: Chain embed_cubin operations for each cubin
        current_obj = unified_obj
        for cubin_name in sorted(embed_cubin.keys()):
            # Create next object file name
            next_obj = f"unified_with_{cubin_name}.o"
            cubin_file = f"{cubin_name}.cubin"

            # Add ninja build rule
            ninja.append(f"build {next_obj}: embed_cubin {current_obj}")
            ninja.append(f"  cubin = {cubin_file}")
            ninja.append(f"  name = {cubin_name}")
            ninja.append("")

            current_obj = next_obj

        # Step 3: Link the final object file
        ninja.append(f"build {output_name}: link {current_obj}")
        ninja.append("")
    else:
        # Directly link object files (for Windows or no cubin embedding)
        link_files_str = " ".join(obj_files)
        ninja.append(f"build {output_name}: link {link_files_str}")
        ninja.append("")

    # default target
    ninja.append(f"default {output_name}")
    ninja.append("")
    return "\n".join(ninja)


def build_ninja(build_dir: str) -> None:
    """Build the module in the given build directory using ninja."""
    command = ["ninja", "-v"]
    num_workers = os.environ.get("MAX_JOBS", None)
    if num_workers is not None:
        command += ["-j", num_workers]
    if IS_WINDOWS:
        status = _run_command_in_dev_prompt(args=command, cwd=build_dir, capture_output=True)
    else:
        status = subprocess.run(check=False, args=command, cwd=build_dir, capture_output=True)
    if status.returncode != 0:
        msg = [f"ninja exited with status {status.returncode}"]
        encoding = "oem" if IS_WINDOWS else "utf-8"
        if status.stdout:
            msg.append(f"stdout:\n{status.stdout.decode(encoding)}")
        if status.stderr:
            msg.append(f"stderr:\n{status.stderr.decode(encoding)}")

        raise RuntimeError("\n".join(msg))


# Translation table for escaping C++ string literals
_CPP_ESCAPE_TABLE = str.maketrans(
    {
        "\\": "\\\\",
        '"': '\\"',
        "\n": "\\n",
        "\r": "\\r",
        "\t": "\\t",
    }
)


def _escape_cpp_string_literal(s: str) -> str:
    """Escape special characters for C++ string literals."""
    return s.translate(_CPP_ESCAPE_TABLE)


def _decorate_with_tvm_ffi(source: str, functions: Mapping[str, str]) -> str:
    """Decorate the given source code with TVM FFI export macros."""
    sources = [
        "#include <tvm/ffi/container/tensor.h>",
        "#include <tvm/ffi/dtype.h>",
        "#include <tvm/ffi/error.h>",
        "#include <tvm/ffi/extra/c_env_api.h>",
        "#include <tvm/ffi/function.h>",
        "",
        source,
    ]

    for func_name, func_doc in functions.items():
        sources.append(f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({func_name}, {func_name});")

        if func_doc:
            # Escape the docstring for C++ string literal
            escaped_doc = _escape_cpp_string_literal(func_doc)
            sources.append(f'TVM_FFI_DLL_EXPORT_TYPED_FUNC_DOC({func_name}, "{escaped_doc}");')

    sources.append("")

    return "\n".join(sources)


def _str_seq2list(seq: Sequence[str] | str | None) -> list[str]:
    if seq is None:
        return []
    elif isinstance(seq, str):
        return [seq]
    else:
        return list(seq)


def _build_impl(  # noqa: PLR0913
    name: str,
    sources: Sequence[str] | str | None,
    extra_cflags: Sequence[str] | None,
    extra_cuda_cflags: Sequence[str] | None,
    extra_ldflags: Sequence[str] | None,
    extra_include_paths: Sequence[str] | None,
    build_directory: str | None,
    need_lock: bool = True,
    embed_cubin: Mapping[str, bytes] | None = None,
    backend: str | None = None,
    output: str | None = None,
) -> str:
    """Real implementation of build function."""
    # need to resolve the path to make it unique
    source_path_list = [str(Path(p).resolve()) for p in _str_seq2list(sources)]
    assert source_path_list, "sources must be provided."

    has_cuda = any(Path(p).suffix.lower() == ".cu" for p in source_path_list)
    resolved_backend = _resolve_gpu_backend(backend) if has_cuda else None
    extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else []
    extra_cflags_list = list(extra_cflags) if extra_cflags is not None else []
    extra_cuda_cflags_list = list(extra_cuda_cflags) if extra_cuda_cflags is not None else []
    extra_include_paths_list = list(extra_include_paths) if extra_include_paths is not None else []

    build_dir: Path
    if build_directory is None:
        cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", str(Path("~/.cache/tvm-ffi").expanduser()))
        source_hash: str = _hash_sources(
            None,
            None,
            source_path_list,
            None,
            {},
            extra_cflags_list,
            extra_cuda_cflags_list,
            extra_ldflags_list,
            extra_include_paths_list,
            embed_cubin,
        )
        build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}"
    else:
        build_dir = Path(build_directory).resolve()
    build_dir.mkdir(parents=True, exist_ok=True)

    # CUBIN embedding is only supported on Unix systems
    if embed_cubin and IS_WINDOWS:
        raise NotImplementedError("CUBIN embedding is not yet supported on Windows")

    # Write CUBIN files to build directory if needed (for Unix systems)
    # These will be embedded using the embed_cubin utility during ninja build
    if embed_cubin:
        for cubin_name, cubin_bytes in embed_cubin.items():
            cubin_path = build_dir / f"{cubin_name}.cubin"
            cubin_path.write_bytes(cubin_bytes)

    # generate build.ninja
    ninja_source = _generate_ninja_build(
        name=name,
        extra_cflags=extra_cflags_list,
        extra_cuda_cflags=extra_cuda_cflags_list,
        extra_ldflags=extra_ldflags_list,
        extra_include_paths=extra_include_paths_list,
        sources=source_path_list,
        embed_cubin=embed_cubin,
        backend=resolved_backend,
        output=output,
    )

    # may not hold lock when build_directory is specified, prevent deadlock
    with FileLock(str(build_dir / "lock")) if need_lock else nullcontext():
        # write build.ninja if it does not already exist
        _maybe_write(str(build_dir / "build.ninja"), ninja_source)
        # build the module
        build_ninja(str(build_dir))
        # Determine the output filename (mirrors _generate_ninja_build logic)
        if output is not None:
            out_ext = Path(output).suffix.lower()
            object_mode = out_ext in (".o", ".obj")
            output_name = Path(output).name
        else:
            object_mode = False
            output_name = f"{name}{'.dll' if IS_WINDOWS else '.so'}"
        if object_mode and IS_WINDOWS:
            # Windows has no ld -r; the actual target is the first intermediate object.
            # The name must match _generate_ninja_build: c_0.o / cpp_0.o / cuda_0.o.
            first_ext = Path(sorted(source_path_list)[0]).suffix.lower() if source_path_list else ""
            if first_ext == ".c":
                obj_name = "c_0.o"
            elif first_ext == ".cu":
                obj_name = "cuda_0.o"
            else:
                obj_name = "cpp_0.o"
            return str((build_dir / obj_name).resolve())
        return str((build_dir / output_name).resolve())


def build_inline(  # noqa: PLR0913
    name: str,
    *,
    cpp_sources: Sequence[str] | str | None = None,
    cuda_sources: Sequence[str] | str | None = None,
    functions: Mapping[str, str] | Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    embed_cubin: Mapping[str, bytes] | None = None,
    backend: str | None = None,
    output: str | None = None,
) -> str:
    """Compile and build a C++/CUDA module from inline source code.

    This function compiles the given C++ and/or CUDA source code into a shared library or object file.
    Both ``cpp_sources`` and ``cuda_sources`` are compiled to an object file. When ``output`` is
    ``None`` (the default) or has a shared-library extension (``.so``, ``.dll``), object files are
    linked into a shared library. When ``output`` has an object-file extension (``.o``, ``.obj``),
    linking is skipped and the path to the object file is returned directly.

    The ``functions`` parameter is used to specify which functions in the source code should be exported to the tvm ffi
    module. It can be a mapping, a sequence, or a single string. When a mapping is given, the keys are the names of the
    exported functions, and the values are docstrings for the functions. When a sequence of string is given, they are
    the function names needed to be exported, and the docstrings are set to empty strings. A single function name can
    also be given as a string, indicating that only one function is to be exported.

    Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags``
    parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional
    flags for your specific use case.

    The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can
    include any header from tvm ffi in your source code. You can also provide additional include paths via the
    ``extra_include_paths`` parameter and include custom headers in your source code.

    The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter
    is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used.
    The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified,
    the default cache directory is ``~/.cache/tvm-ffi``.

    Parameters
    ----------
    name
        The name of the tvm ffi module.
    cpp_sources
        The C++ source code. It can be a list of sources or a single source.
    cuda_sources
        The CUDA source code. It can be a list of sources or a single source.
    functions
        The functions in cpp_sources or cuda_source that will be exported to the tvm ffi module. When a mapping is
        given, the keys are the names of the exported functions, and the values are docstrings for the functions
        (use an empty string to skip documentation for specific functions). When a sequence or a single string is given, they are
        the functions needed to be exported, and the docstrings are set to empty strings. A single function name can
        also be given as a string. When cpp_sources is given, the functions must be declared (not necessarily defined)
        in the cpp_sources. When cpp_sources is not given, the functions must be defined in the cuda_sources. If not
        specified, no function will be exported.
    extra_cflags
        The extra compiler flags for C++ compilation.
        The default flags are:

        - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
        - On Windows: ['/std:c++17', '/O2']

    extra_cuda_cflags
        The extra compiler flags for CUDA compilation.

    extra_ldflags
        The extra linker flags.
        The default flags are:

        - On Linux/macOS: ['-shared']
        - On Windows: ['/DLL']

    extra_include_paths
        The extra include paths.

    build_directory
        The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the
        cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to
        specify the cache directory.

    embed_cubin: Mapping[str, bytes], optional
        A mapping from CUBIN module names to CUBIN binary data. TVM-FFI provides a macro `TVM_FFI_EMBED_CUBIN(name)` to embed
        CUBIN data into the compiled shared library. The keys should match the names used in `TVM_FFI_EMBED_CUBIN(name)` calls
        in the C++ source code. The values are the CUBIN binary data bytes. The embedded CUBIN kernels can be accessed by
        the macro `TVM_FFI_EMBED_CUBIN_GET_KERNEL(name, kernel_name)` defined in the `tvm/ffi/extra/cuda/cubin_launcher.h` header.
        See the `examples/cubin_launcher` directory for examples how to use cubin launcher to launch CUBIN kernels in TVM-FFI.

    backend
        The GPU backend to use. It can be "cuda" or "hip".
        If not specified, the backend will be automatically determined based on the available GPU and the provided source code.

    output
        Output filename that determines the build type from its extension. When ``None``
        (the default), builds a shared library (``.so`` on Unix, ``.dll`` on Windows).
        Use an object-file extension (e.g., ``"hello.o"``) to skip linking and produce
        a relocatable object file. The file is placed in the build directory.

    Returns
    -------
    path: str
        The path to the built shared library or object file.

    Example
    -------

    .. code-block:: python

        import torch
        from tvm_ffi import Module
        import tvm_ffi.cpp

        # define the cpp source code
        cpp_source = '''
             void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
               // implementation of a library function
               TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor";
               DLDataType f32_dtype{kDLFloat, 32, 1};
               TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor";
               TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor";
               TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor";
               TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape";
               for (int i = 0; i < x.size(0); ++i) {
                 static_cast<float*>(y.data_ptr())[i] = static_cast<float*>(x.data_ptr())[i] + 1;
               }
             }
        '''

        # compile the cpp source code and load the module
        lib_path: str = tvm_ffi.cpp.build_inline(
            name="hello",
            cpp_sources=cpp_source,
            functions="add_one_cpu",
        )

        # load the module
        mod: Module = tvm_ffi.load_module(lib_path)

        # use the function from the loaded module to perform
        x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
        y = torch.empty_like(x)
        mod.add_one_cpu(x, y)
        torch.testing.assert_close(x + 1, y)

    """
    cpp_source_list = _str_seq2list(cpp_sources)
    cpp_source = "\n".join(cpp_source_list)
    with_cpp = bool(cpp_source_list)
    del cpp_source_list

    cuda_source_list = _str_seq2list(cuda_sources)
    cuda_source = "\n".join(cuda_source_list)
    with_backend = bool(cuda_source_list)
    del cuda_source_list

    extra_ldflags_list = list(extra_ldflags) if extra_ldflags is not None else []
    extra_cflags_list = list(extra_cflags) if extra_cflags is not None else []
    extra_cuda_cflags_list = list(extra_cuda_cflags) if extra_cuda_cflags is not None else []
    extra_include_paths_list = list(extra_include_paths) if extra_include_paths is not None else []

    # add function registration code to sources
    if functions is None:
        function_map: dict[str, str] = {}
    elif isinstance(functions, str):
        function_map = {functions: ""}
    elif isinstance(functions, Mapping):
        function_map = dict(functions)
    else:
        function_map = {name: "" for name in functions}

    if with_cpp:
        cpp_source = _decorate_with_tvm_ffi(cpp_source, function_map)
        cuda_source = _decorate_with_tvm_ffi(cuda_source, {})
    else:
        cpp_source = _decorate_with_tvm_ffi(cpp_source, {})
        cuda_source = _decorate_with_tvm_ffi(cuda_source, function_map)
    # determine the cache dir for the built module
    build_dir: Path
    if build_directory is None:
        cache_dir = os.environ.get("TVM_FFI_CACHE_DIR", str(Path("~/.cache/tvm-ffi").expanduser()))
        source_hash: str = _hash_sources(
            cpp_source,
            cuda_source,
            None,
            None,
            function_map,
            extra_cflags_list,
            extra_cuda_cflags_list,
            extra_ldflags_list,
            extra_include_paths_list,
            embed_cubin,
        )
        build_dir = Path(cache_dir).expanduser() / f"{name}_{source_hash}"
    else:
        build_dir = Path(build_directory).resolve()
    build_dir.mkdir(parents=True, exist_ok=True)

    cpp_file = str((build_dir / "main.cpp").resolve())
    cuda_file = str((build_dir / "cuda.cu").resolve())

    with FileLock(str(build_dir / "lock")):
        # write source files if they do not already exist
        _maybe_write(cpp_file, cpp_source)
        if with_backend:
            _maybe_write(cuda_file, cuda_source)

        src_files = []
        if with_cpp:
            src_files.append(cpp_file)
        if with_backend:
            src_files.append(cuda_file)
        return _build_impl(
            name=name,
            sources=src_files,
            extra_cflags=extra_cflags_list,
            extra_cuda_cflags=extra_cuda_cflags_list,
            extra_ldflags=extra_ldflags_list,
            extra_include_paths=extra_include_paths_list,
            build_directory=str(build_dir),
            need_lock=False,  # already hold the lock
            embed_cubin=embed_cubin,
            backend=backend,
            output=output,
        )


def load_inline(  # noqa: PLR0913
    name: str,
    *,
    cpp_sources: Sequence[str] | str | None = None,
    cuda_sources: Sequence[str] | str | None = None,
    functions: Mapping[str, str] | Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    embed_cubin: Mapping[str, bytes] | None = None,
    keep_module_alive: bool = True,
    backend: str | None = None,
) -> Module:
    """Compile, build and load a C++/CUDA module from inline source code.

    This function compiles the given C++ and/or CUDA source code into a shared library. Both ``cpp_sources`` and
    ``cuda_sources`` are compiled to an object file, and then linked together into a shared library. It's possible to only
    provide cpp_sources or cuda_sources.

    The ``functions`` parameter is used to specify which functions in the source code should be exported to the tvm ffi
    module. It can be a mapping, a sequence, or a single string. When a mapping is given, the keys are the names of the
    exported functions, and the values are docstrings for the functions. When a sequence of string is given, they are
    the function names needed to be exported, and the docstrings are set to empty strings. A single function name can
    also be given as a string, indicating that only one function is to be exported.

    Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags``
    parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional
    flags for your specific use case.

    The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can
    include any header from tvm ffi in your source code. You can also provide additional include paths via the
    ``extra_include_paths`` parameter and include custom headers in your source code.

    The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter
    is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used.
    The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified,
    the default cache directory is ``~/.cache/tvm-ffi``.

    Parameters
    ----------
    name
        The name of the tvm ffi module.
    cpp_sources
        The C++ source code. It can be a list of sources or a single source.
    cuda_sources
        The CUDA source code. It can be a list of sources or a single source.
    functions
        The functions in cpp_sources or cuda_source that will be exported to the tvm ffi module. When a mapping is
        given, the keys are the names of the exported functions, and the values are docstrings for the functions
        (use an empty string to skip documentation for specific functions). When a sequence or a single string is given, they are
        the functions needed to be exported, and the docstrings are set to empty strings. A single function name can
        also be given as a string. When cpp_sources is given, the functions must be declared (not necessarily defined)
        in the cpp_sources. When cpp_sources is not given, the functions must be defined in the cuda_sources. If not
        specified, no function will be exported.
    extra_cflags
        The extra compiler flags for C++ compilation.
        The default flags are:

        - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
        - On Windows: ['/std:c++17', '/O2']

    extra_cuda_cflags
        The extra compiler flags for CUDA compilation.

    extra_ldflags
        The extra linker flags.
        The default flags are:

        - On Linux/macOS: ['-shared']
        - On Windows: ['/DLL']

    extra_include_paths
        The extra include paths.

    build_directory
        The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the
        cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to
        specify the cache directory.

    embed_cubin
        A mapping from CUBIN module names to CUBIN binary data. When provided, the CUBIN data will be embedded
        into the compiled shared library using objcopy, making it accessible via the TVM_FFI_EMBED_CUBIN macro.
        The keys should match the names used in TVM_FFI_EMBED_CUBIN calls in the C++ source code.

    keep_module_alive
        Whether to keep the module alive. If True, the module will be kept alive
        for the duration of the program until libtvm_ffi.so is unloaded.

    backend
        The GPU backend to use. It can be "cuda" or "hip".
        If not specified, the backend will be automatically determined based on the available GPU and the provided source code.

    Returns
    -------
    mod: Module
        The loaded tvm ffi module.

    See Also
    --------
    :py:func:`tvm_ffi.load_module`

    Example
    -------

    .. code-block:: python

        import torch
        from tvm_ffi import Module
        import tvm_ffi.cpp

        # define the cpp source code
        cpp_source = '''
             void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
               // implementation of a library function
               TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor";
               DLDataType f32_dtype{kDLFloat, 32, 1};
               TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor";
               TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor";
               TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor";
               TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape";
               for (int i = 0; i < x.size(0); ++i) {
                 static_cast<float*>(y.data_ptr())[i] = static_cast<float*>(x.data_ptr())[i] + 1;
               }
             }
        '''

        # compile the cpp source code and load the module
        mod: Module = tvm_ffi.cpp.load_inline(
            name="hello",
            cpp_sources=cpp_source,
            functions="add_one_cpu",
        )

        # use the function from the loaded module to perform
        x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
        y = torch.empty_like(x)
        mod.add_one_cpu(x, y)
        torch.testing.assert_close(x + 1, y)

    """
    return load_module(
        build_inline(
            name=name,
            cpp_sources=cpp_sources,
            cuda_sources=cuda_sources,
            functions=functions,
            extra_cflags=extra_cflags,
            extra_cuda_cflags=extra_cuda_cflags,
            extra_ldflags=extra_ldflags,
            extra_include_paths=extra_include_paths,
            build_directory=build_directory,
            embed_cubin=embed_cubin,
            backend=backend,
        ),
        keep_module_alive=keep_module_alive,
    )


def build(  # noqa: PLR0913
    name: str,
    *,
    sources: Sequence[str] | str | None = None,
    cpp_files: Sequence[str] | str | None = None,
    cuda_files: Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    backend: str | None = None,
    output: str | None = None,
) -> str:
    """Compile and build a C/C++/CUDA module from source files.

    This function compiles the given C, C++, and/or CUDA source files into a shared library or
    object file. The compiler is selected automatically based on file extension:

    - ``.c`` — compiled with the C compiler (``$CC``)
    - ``.cc``, ``.cpp``, ``.cxx`` — compiled with the C++ compiler (``$CXX``)
    - ``.o``, ``.obj`` — pre-compiled objects, passed directly to the linker

    When ``output`` is ``None`` (the default) or has a shared-library extension, object files are
    linked into a shared library. When ``output`` has an object-file extension (``.o``, ``.obj``),
    linking is skipped and the path to the object file is returned.

    Note that this function does not automatically export functions to the tvm ffi module. You need to
    manually use the TVM FFI export macros (e.g., ``TVM_FFI_DLL_EXPORT_TYPED_FUNC``) in your source files to export
    functions. This gives you more control over which functions are exported and how they are exported.

    Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags``
    parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional
    flags for your specific use case.

    The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can
    include any header from tvm ffi in your source files. You can also provide additional include paths via the
    ``extra_include_paths`` parameter and include custom headers in your source code.

    The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter
    is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used.
    The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified,
    the default cache directory is ``~/.cache/tvm-ffi``.

    The C compiler is controlled by the ``$CC`` environment variable (default: ``cc`` on Unix, ``cl`` on Windows).
    The C++ compiler is controlled by the ``$CXX`` environment variable (default: ``c++`` on Unix, ``cl`` on Windows).

    Parameters
    ----------
    name
        The name of the tvm ffi module.
    sources
        Source files to compile. The compiler is auto-detected from the file extension:

        - ``.c`` → C compiler (``$CC``)
        - ``.cc``, ``.cpp``, ``.cxx`` → C++ compiler (``$CXX``)
        - ``.cu`` → CUDA/HIP compiler (``nvcc`` or ``hipcc``)
        - ``.o``, ``.obj`` → pre-compiled objects, passed directly to the linker

        It can be a list of file paths or a single file path.
    cpp_files
        Alias for ``sources``, kept for backward compatibility.
    cuda_files
        Alias for ``sources``, kept for backward compatibility.
    extra_cflags
        Extra compiler flags applied to both C and C++ compilation.
        The C++ default flags are:

        - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
        - On Windows: ['/std:c++17', '/MD', '/O2']

        The C default flags omit ``-std=c++17`` and ``/EHsc``.

    extra_cuda_cflags
        The extra compiler flags for CUDA compilation.
        The default flags are:

        - ['-Xcompiler', '-fPIC', '-std=c++17', '-O2'] (Linux/macOS)
        - ['-Xcompiler', '/std:c++17', '/O2'] (Windows)

    extra_ldflags
        The extra linker flags.
        The default flags are:

        - On Linux/macOS: ['-shared', '-L<tvm_ffi_lib_path>', '-ltvm_ffi']
        - On Windows: ['/DLL', '/LIBPATH:<tvm_ffi_lib_path>', '<tvm_ffi_lib_name>.lib']

    extra_include_paths
        The extra include paths for header files. Both absolute and relative paths are supported.

    build_directory
        The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the
        cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to
        specify the cache directory.

    backend
        The GPU backend to use. It can be "cuda" or "hip".
        If not specified, the backend will be automatically determined based on the available GPU and the provided source code.

    output
        Output filename that determines the build type from its extension. When ``None``
        (the default), builds a shared library (``.so`` on Unix, ``.dll`` on Windows).
        Use an object-file extension (e.g., ``"my_ops.o"``) to skip linking and produce
        a relocatable object file. The file is placed in the build directory.

    Returns
    -------
    path: str
        The path to the built shared library or object file.

    Example
    -------

    .. code-block:: python

        import torch
        from tvm_ffi import Module
        import tvm_ffi.cpp

        # Assume we have a C++ source file "my_ops.cpp" with the following content:
        # ```cpp
        # #include <tvm/ffi/container/tensor.h>
        # #include <tvm/ffi/dtype.h>
        # #include <tvm/ffi/error.h>
        # #include <tvm/ffi/extra/c_env_api.h>
        # #include <tvm/ffi/function.h>
        #
        # void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
        #   TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor";
        #   DLDataType f32_dtype{kDLFloat, 32, 1};
        #   TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor";
        #   TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor";
        #   TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor";
        #   TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape";
        #   for (int i = 0; i < x.size(0); ++i) {
        #     static_cast<float*>(y.data_ptr())[i] = static_cast<float*>(x.data_ptr())[i] + 1;
        #   }
        # }
        #
        # TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu);
        # ```

        # compile the cpp source file and get the library path
        lib_path: str = tvm_ffi.cpp.build(
            name="my_ops",
            sources="my_ops.cpp",
        )

        # load the module
        mod: Module = tvm_ffi.load_module(lib_path)

        # use the function from the loaded module
        x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
        y = torch.empty_like(x)
        mod.add_one_cpu(x, y)
        torch.testing.assert_close(x + 1, y)

    """
    # Merge sources, cpp_files, and cuda_files (backward compat aliases)
    merged = _str_seq2list(sources) + _str_seq2list(cpp_files) + _str_seq2list(cuda_files)
    return _build_impl(
        name=name,
        sources=merged or None,
        extra_cflags=extra_cflags,
        extra_cuda_cflags=extra_cuda_cflags,
        extra_ldflags=extra_ldflags,
        extra_include_paths=extra_include_paths,
        build_directory=build_directory,
        need_lock=True,
        backend=backend,
        output=output,
    )


def load(  # noqa: PLR0913
    name: str,
    *,
    sources: Sequence[str] | str | None = None,
    cpp_files: Sequence[str] | str | None = None,
    cuda_files: Sequence[str] | str | None = None,
    extra_cflags: Sequence[str] | None = None,
    extra_cuda_cflags: Sequence[str] | None = None,
    extra_ldflags: Sequence[str] | None = None,
    extra_include_paths: Sequence[str] | None = None,
    build_directory: str | None = None,
    keep_module_alive: bool = True,
    backend: str | None = None,
) -> Module:
    """Compile, build and load a C/C++/CUDA module from source files.

    This function compiles the given source files into a shared library and loads it as a tvm ffi
    module. The compiler is selected automatically based on file extension.

    Note that this function does not automatically export functions to the tvm ffi module. You need to
    manually use the TVM FFI export macros (e.g., :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC`) in your source files to export
    functions. This gives you more control over which functions are exported and how they are exported.

    Extra compiler and linker flags can be provided via the ``extra_cflags``, ``extra_cuda_cflags``, and ``extra_ldflags``
    parameters. The default flags are generally sufficient for most use cases, but you may need to provide additional
    flags for your specific use case.

    The include dir of tvm ffi and dlpack are used by default for the compiler to find the headers. Thus, you can
    include any header from tvm ffi in your source files. You can also provide additional include paths via the
    ``extra_include_paths`` parameter and include custom headers in your source code.

    The compiled shared library is cached in a cache directory to avoid recompilation. The `build_directory` parameter
    is provided to specify the build directory. If not specified, a default tvm ffi cache directory will be used.
    The default cache directory can be specified via the `TVM_FFI_CACHE_DIR` environment variable. If not specified,
    the default cache directory is ``~/.cache/tvm-ffi``.

    Parameters
    ----------
    name
        The name of the tvm ffi module.
    sources
        Source files to compile. The compiler is auto-detected from the file extension:
        ``.c`` → C, ``.cc``/``.cpp``/``.cxx`` → C++, ``.cu`` → CUDA/HIP,
        ``.o``/``.obj`` → linker passthrough. It can be a list of file paths or a single file path.
    cpp_files
        Alias for ``sources``, kept for backward compatibility.
    cuda_files
        Alias for ``sources``, kept for backward compatibility.
    extra_cflags
        The extra compiler flags for C++ compilation.
        The default flags are:

        - On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
        - On Windows: ['/std:c++17', '/MD', '/O2']

    extra_cuda_cflags
        The extra compiler flags for CUDA compilation.
        The default flags are:

        - ['-Xcompiler', '-fPIC', '-std=c++17', '-O2'] (Linux/macOS)
        - ['-Xcompiler', '/std:c++17', '/O2'] (Windows)

    extra_ldflags
        The extra linker flags.
        The default flags are:

        - On Linux/macOS: ['-shared', '-L<tvm_ffi_lib_path>', '-ltvm_ffi']
        - On Windows: ['/DLL', '/LIBPATH:<tvm_ffi_lib_path>', '<tvm_ffi_lib_name>.lib']

    extra_include_paths
        The extra include paths for header files. Both absolute and relative paths are supported.

    build_directory
        The build directory. If not specified, a default tvm ffi cache directory will be used. By default, the
        cache directory is ``~/.cache/tvm-ffi``. You can also set the ``TVM_FFI_CACHE_DIR`` environment variable to
        specify the cache directory.

    keep_module_alive
        Whether to keep the module alive. If True, the module will be kept alive
        for the duration of the program until libtvm_ffi.so is unloaded.

    backend
        The GPU backend to use. It can be "cuda" or "hip".
        If not specified, the backend will be automatically determined based on the available GPU and the provided source code.

    Returns
    -------
    mod: Module
        The loaded tvm ffi module.

    See Also
    --------
    :py:func:`tvm_ffi.load_module`

    Example
    -------

    .. code-block:: python

        import torch
        from tvm_ffi import Module
        import tvm_ffi.cpp

        # Assume we have a C++ source file "my_ops.cpp" with the following content:
        # ```cpp
        # #include <tvm/ffi/container/tensor.h>
        # #include <tvm/ffi/dtype.h>
        # #include <tvm/ffi/error.h>
        # #include <tvm/ffi/extra/c_env_api.h>
        # #include <tvm/ffi/function.h>
        #
        # void add_one_cpu(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
        #   TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor";
        #   DLDataType f32_dtype{kDLFloat, 32, 1};
        #   TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor";
        #   TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor";
        #   TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor";
        #   TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape";
        #   for (int i = 0; i < x.size(0); ++i) {
        #     static_cast<float*>(y.data_ptr())[i] = static_cast<float*>(x.data_ptr())[i] + 1;
        #   }
        # }
        #
        # TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, add_one_cpu);
        # ```

        # compile the cpp source file and load the module
        mod: Module = tvm_ffi.cpp.load(
            name="my_ops",
            sources="my_ops.cpp",
        )

        # use the function from the loaded module
        x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
        y = torch.empty_like(x)
        mod.add_one_cpu(x, y)
        torch.testing.assert_close(x + 1, y)

    """
    return load_module(
        build(
            name=name,
            sources=sources,
            cpp_files=cpp_files,
            cuda_files=cuda_files,
            extra_cflags=extra_cflags,
            extra_cuda_cflags=extra_cuda_cflags,
            extra_ldflags=extra_ldflags,
            extra_include_paths=extra_include_paths,
            build_directory=build_directory,
            backend=backend,
        ),
        keep_module_alive=keep_module_alive,
    )
