#!/usr/bin/env python3
"""Probe the installed toolkit instead of maintaining a CUDA/compiler version table."""
from __future__ import annotations

import argparse
import os
from pathlib import Path
import shlex
import subprocess
import tempfile


def check_toolchain(nvcc: str, cc: str, cxx: str, *, standard: str = '17',
                    c_args: list[str] | None = None,
                    cxx_args: list[str] | None = None) -> str:
    version = subprocess.run([nvcc, '--version'], text=True, capture_output=True, check=True).stdout.strip()
    # Do not permit ambient flags to bypass the installed toolkit's version checks.
    env = dict(os.environ)
    for key in ('NVCC_PREPEND_FLAGS', 'NVCC_APPEND_FLAGS', 'NVCC_CCBIN'):
        env.pop(key, None)
    with tempfile.TemporaryDirectory(prefix='mfem-cuda-check-') as directory:
        root = Path(directory)
        source = root / 'probe.cu'
        source.write_text('''#include <cuda_runtime.h>
#include <thrust/device_vector.h>
#include <cmath>
__global__ void probe(int *p) { *p = 1; }
int main() { int *p = nullptr; cudaMalloc(&p, sizeof(int));
probe<<<1, 1>>>(p); return (int)cudaFree(p); }
''')
        for label, host, flags in [('C', cc, c_args or []), ('C++', cxx, cxx_args or [])]:
            command = [nvcc, '--compiler-bindir', host, '--std=c++' + standard,
                       str(source), '-o', str(root / ('probe.exe' if os.name == 'nt' else 'probe'))]
            if label == 'C':
                command.append('-c')
            for flag in flags:
                command += ['-Xcompiler', flag]
            result = subprocess.run(command, env=env, text=True, capture_output=True)
            if result.returncode:
                raise RuntimeError(
                    f'CUDA compatibility check failed for the selected {label} compiler: {host}\n'
                    f'{version}\nCommand: {shlex.join(command)}\n'
                    f'{result.stdout}{result.stderr}\n'
                    'Select a C/C++ toolchain supported by this CUDA toolkit using CC and CXX '
                    '(for example CC=gcc-15 CXX=g++-15), or install a compatible toolkit. '
                    'Use a fresh Meson build directory when changing compilers; --reconfigure '
                    'does not change cached compilers. For a CPU build use -Dmfem_cuda=disabled. '
                    'The probe compiles and links but does not require a running GPU.'
                )
    return f'CUDA compatibility checks passed: C={cc}, C++={cxx}, CUDA C++{standard}\n{version}'


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--nvcc', required=True)
    parser.add_argument('--cc', required=True)
    parser.add_argument('--cxx', required=True)
    parser.add_argument('--standard', default='17')
    parser.add_argument('--c-arg', action='append', default=[])
    parser.add_argument('--cxx-arg', action='append', default=[])
    args = parser.parse_args()
    try:
        print(check_toolchain(args.nvcc, args.cc, args.cxx, standard=args.standard,
                              c_args=args.c_arg, cxx_args=args.cxx_arg))
    except (RuntimeError, OSError, subprocess.SubprocessError) as error:
        print(error)
        return 1
    return 0


if __name__ == '__main__':
    raise SystemExit(main())