import os, sys, shutil
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext


class ObjCppBuildExt(build_ext):
    def build_extensions(self):
        import torch
        from torch.utils import cpp_extension

        self.compiler.src_extensions.append('.mm')
        original_compile = self.compiler._compile

        def objcpp_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
            if src.endswith('.mm'):
                extra_postargs = ['-x', 'objective-c++'] + list(extra_postargs or [])
            return original_compile(obj, src, ext, cc_args, extra_postargs, pp_opts)

        self.compiler._compile = objcpp_compile

        for ext in self.extensions:
            ext.include_dirs.extend(cpp_extension.include_paths())
            ext.library_dirs.extend(cpp_extension.library_paths())
            ext.libraries.extend(['c10', 'torch', 'torch_cpu', 'torch_python'])

        super().build_extensions()


def get_extensions():
    if sys.platform != 'darwin':
        return []
    return [Extension(
        name='mps_linear_attn._C',
        sources=['mps_linear_attn/csrc/mla.mm'],
        extra_compile_args=['-std=c++17', '-O3', '-DTORCH_EXTENSION_NAME=_C'],
        extra_link_args=['-framework', 'Metal', '-framework', 'Foundation'],
    )]


setup(
    name='mps-linear-attn',
    version='0.1.0',
    packages=find_packages(),
    install_requires=['torch>=2.0'],
    ext_modules=get_extensions(),
    cmdclass={'build_ext': ObjCppBuildExt},
)
