from setuptools import setup, find_packages
import os
import platform
from torch.utils.cpp_extension import BuildExtension, CppExtension

CSRC_DIR = os.path.join('toploc', 'C', 'csrc')

# Define mac-specific compiler and linker flags
extra_compile_args = ['-O3']
extra_link_args: list[str] = []

# Add macOS specific flags
if platform.system() == 'Darwin':
    # Enable support for both Intel and Apple Silicon
    extra_compile_args.extend(['-arch', 'x86_64', '-arch', 'arm64'])
    extra_link_args.extend(['-arch', 'x86_64', '-arch', 'arm64'])
    
    # Add minimum deployment target for macOS
    extra_compile_args.append('-mmacosx-version-min=10.13')
    extra_link_args.append('-mmacosx-version-min=10.13')

extensions = [
    CppExtension(
        name='toploc.C.csrc.ndd',
        sources=[os.path.join(CSRC_DIR, 'ndd.cpp')],
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    ),
    CppExtension(
        name='toploc.C.csrc.utils',
        sources=[os.path.join(CSRC_DIR, 'utils.cpp')],
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    ),
]

setup(
    name='toploc',
    ext_modules=extensions,
    packages=['toploc', 'toploc.C.csrc'],
    package_data={
        'toploc.C.csrc': ['*.pyi'],  # Include .pyi files
    },
    cmdclass={'build_ext': BuildExtension},
)
