"""Make the CUDA build of llama-cpp-python load on Windows.

The CUDA wheel ships `ggml-cuda.dll`, which depends on the CUDA runtime
(`cudart64_12.dll`) and cuBLAS (`cublas64_12.dll`, `cublasLt64_12.dll`). Those
are NOT bundled; we install them as pip wheels (`nvidia-cuda-runtime-cu12`,
`nvidia-cublas-cu12`), which drop the DLLs under `site-packages/nvidia/*/bin`.

llama-cpp-python loads its shared library with `winmode=RTLD_GLOBAL`, and under
that flag Windows ignores both `PATH` and `os.add_dll_directory` when resolving
dependencies — so simply registering the nvidia dirs isn't enough. The reliable
fix is to *pre-load* the DLLs ourselves with the default loader (which does
honor `add_dll_directory`) in dependency order. Once a module is mapped into the
process, llama-cpp-python's later load reuses the in-memory copy instead of
searching the disk again.

No-op on non-Windows and when the CUDA/nvidia wheels aren't present (CPU build).
"""

import ctypes
import glob
import os
import sys

_done = False

# llama.dll last — it depends on all the others. ggml-cuda.dll is the one that
# pulls in the CUDA runtime; the rest are loaded so their inter-deps resolve.
_PRELOAD_ORDER = [
    "ggml-base.dll",
    "ggml.dll",
    "ggml-cpu.dll",
    "ggml-cuda.dll",
    "llama.dll",
]


def _site_packages_dirs():
    seen = set()
    for p in sys.path:
        if p.endswith("site-packages") and p not in seen:
            seen.add(p)
            yield p
    # Fallback: the repo-local venv, in case this runs outside the venv.
    local = os.path.join(os.path.dirname(__file__), ".venv", "Lib", "site-packages")
    if os.path.isdir(local) and local not in seen:
        yield local


def ensure() -> None:
    global _done
    if _done:
        return
    _done = True
    if sys.platform != "win32":
        return

    lib_dir = None
    for sp in _site_packages_dirs():
        nvidia_root = os.path.join(sp, "nvidia")
        if os.path.isdir(nvidia_root):
            for bin_dir in glob.glob(os.path.join(nvidia_root, "*", "bin")):
                try:
                    os.add_dll_directory(bin_dir)
                except OSError:
                    pass
        candidate = os.path.join(sp, "llama_cpp", "lib")
        if os.path.isdir(candidate):
            lib_dir = candidate

    if not lib_dir or not os.path.exists(os.path.join(lib_dir, "ggml-cuda.dll")):
        return  # CPU build (no CUDA dll) — nothing to pre-load

    try:
        os.add_dll_directory(lib_dir)
    except OSError:
        pass
    for name in _PRELOAD_ORDER:
        path = os.path.join(lib_dir, name)
        if os.path.exists(path):
            try:
                ctypes.CDLL(path)
            except OSError:
                # Leave it to llama-cpp-python to surface a precise error.
                pass