from setuptools import setup, Extension, find_packages
import numpy
import os
import sys
import subprocess

def detect_nvc():
    """
    Detect if the NVIDIA HPC Compiler (nvc) is being used.
    This compiler is often default on HPC modules but fails to build
    standard C-extensions expecting GCC behavior.
    """
    # 1. Check environment variable
    cc = os.environ.get("CC", "")
    if "nvc" in cc:
        return True

    # 2. Check default 'cc' command if CC is not explicitly set
    # Only run this check if CC is not set, to avoid overriding user's explicit choice.
    if not cc:
        try:
            # Run 'cc --version' to see who provides it
            out = subprocess.check_output(
                ["cc", "--version"], stderr=subprocess.STDOUT
            ).decode().lower()
            return "nvc" in out or "nvidia" in out
        except Exception:
            # If 'cc' command fails or doesn't exist, assume safe (or failure will happen later)
            return False
    return False

if detect_nvc():
    sys.exit(
        "\n" + "=" * 60 + "\n"
        " ERROR: NVIDIA nvc compiler detected!\n"
        "\n"
        " The 'nvc' compiler is not compatible with some C-extensions in this package.\n"
        " Please switch to GCC (GNU Compiler Collection) before installing.\n"
        "\n"
        " Quick Fix for HPC Users:\n"
        "   export CC=gcc\n"
        "   export CXX=g++\n"
        "   export FC=gfortran\n"
        "   pip install .\n"
        "=" * 60 + "\n"
    )

# Define the C source directory
c_src_dir = os.path.join('macer', 'externals', 'dynaphopy_bundled', 'c')

# Define the extensions
# 1. displacements
ext_displacements = Extension(
    name='macer.externals.dynaphopy_bundled.c.displacements',
    sources=[os.path.join(c_src_dir, 'displacements.c')],
    include_dirs=[numpy.get_include(), c_src_dir],
    extra_compile_args=['-O3', '-fPIC'],
)

# 2. mem
ext_mem = Extension(
    name='macer.externals.dynaphopy_bundled.c.mem',
    sources=[os.path.join(c_src_dir, 'mem.c'),
             os.path.join(c_src_dir, 'displacements.c')],
    include_dirs=[numpy.get_include(), c_src_dir],
    extra_compile_args=['-O3', '-fPIC'],
)

# 3. correlation (Added)
ext_correlation = Extension(
    name='macer.externals.dynaphopy_bundled.c.correlation',
    sources=[os.path.join(c_src_dir, 'correlation.c')],
    include_dirs=[numpy.get_include(), c_src_dir],
    extra_compile_args=['-O3', '-fPIC'],
)

ext_modules = [ext_displacements, ext_mem, ext_correlation]
if sys.platform.startswith("win"):
    # MSVC fails to compile the bundled C sources (C99 complex).
    # Allow Windows wheels to build without these optional extensions.
    ext_modules = []

setup(
    packages=find_packages(),
    ext_modules=ext_modules,
)
