import os
import shutil
import subprocess
import glob
import torch
import torch_npu
import stat
import sysconfig

from setuptools import setup, Extension, Command
from pybind11.setup_helpers import Pybind11Extension, build_ext
from setuptools.command.build_clib import build_clib
from setuptools.command.install import install
from distutils.command.build_py import build_py

from distutils.command.clean import clean as _clean

BASE_DIR = os.path.dirname(os.path.realpath(__file__))
BUILD_DIR = os.path.join(BASE_DIR, "build")
if os.path.exists(BUILD_DIR):
    shutil.rmtree(BUILD_DIR)

class clean(_clean):
    def run(self):
        super().run()
        
        build_dirs = ['build', 'dist']
        for d in build_dirs:
            dir_path = os.path.join(os.getcwd(), d)
            if os.path.exists(dir_path):
                print(f"Cleaning up {dir_path}...")
                shutil.rmtree(dir_path)

class CustomUninstallCommand(Command):
    """Custom command to uninstall the package and clean up."""
    user_options = []

    def run(self):
        subprocess.run(['pip', 'uninstall', '-y', self.distribution.get_name()])
        self.distribution.run_command('clean')


asc_path = os.getenv("ASCEND_HOME_PATH", "")
if asc_path == "":
    raise Exception("ASCEND_HOME_PATH is not set, source <ascend-toolkit>/set_env.sh first")


def _build_npu_common():
    build_dir = os.path.join(BASE_DIR, "build")
    lib_dir = os.path.join(build_dir, "npu_inductor_mlir/lib")
    os.makedirs(lib_dir, exist_ok=True)
    os.chmod(lib_dir, mode=stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR | stat.S_IRGRP | stat.S_IXGRP)
    src_path = os.path.join(BASE_DIR, "npu_inductor_mlir/cpp_common/cpp_common.cpp")
    so_path = os.path.join(lib_dir, "libcpp_common.so")

    cxx = os.environ.get("CC")
    if cxx is None:
        clangxx = shutil.which("clang++")
        gxx = shutil.which("g++")
        cxx = clangxx if clangxx is not None else gxx
        if cxx is None:
            raise RuntimeError("Failed to find C++ compiler")
    cc_cmd = [cxx, src_path]

    # find the python library
    if hasattr(sysconfig, 'get_default_scheme'):
        scheme = sysconfig.get_default_scheme()
    else:
        scheme = sysconfig._get_default_scheme()
    # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install
    # path changes to include 'local'. This change is required to use triton with system-wide python.
    if scheme == 'posix_local':
        scheme = 'posix_prefix'
    py_include_dir = sysconfig.get_paths(scheme=scheme)["include"]

    cc_cmd += [f"-I{py_include_dir}"]

    # find the ascend library
    cc_cmd += [f"-I{os.path.join(BASE_DIR, 'npu_inductor_mlir/_C/include')}"] # hacl_rt.h
    torch_path = os.path.dirname(os.path.realpath(torch.__file__))
    torch_npu_path = os.path.dirname(os.path.realpath(torch_npu.__file__))
    cc_cmd += [
        f"-I{os.path.join(asc_path, 'include')}",
        f"-L{os.path.join(asc_path, 'lib64')}",
        f"-I{os.path.join(torch_path, 'include')}",
        f"-I{os.path.join(torch_npu_path, 'include')}",
        f"-L{os.path.join(torch_npu_path, 'lib')}",
        "-lruntime", "-lascendcl", "-ltorch_npu", "-lprofapi",
        "-std=c++17", "-D_GLIBCXX_USE_CXX11_ABI=0", "-shared"
    ]
    cc_cmd += ["-fPIC", "-o", so_path]
    ret = subprocess.check_call(cc_cmd)

    if ret != 0:
        raise RuntimeError("Failed to compile " + src_path)


class CPPLibBuild(build_clib, object):
    def run(self):
        _build_npu_common()


class Build(build_ext):
    def run(self):
        self.run_command('build_clib')
        self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build"))
        super(Build, self).run()


class PythonPackageBuild(build_py, object):
    def run(self) -> None:
        build_dir = os.path.join(BASE_DIR, "build")
        os.makedirs(build_dir, exist_ok=True)
        src_files = glob.glob(
            os.path.join(BASE_DIR, "npu_inductor_mlir", '*.py'),
            recursive=True) + glob.glob(
            os.path.join(BASE_DIR, "npu_inductor_mlir", '**/*.py'),
            recursive=True) + glob.glob(
            os.path.join(BASE_DIR, "npu_inductor_mlir", '**/*.h'),
            recursive=True)

        for src in src_files:
            dst = os.path.join(
                os.path.join(BASE_DIR, "build/npu_inductor_mlir"),
                os.path.relpath(src, os.path.join(BASE_DIR, "npu_inductor_mlir")))
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            self.copy_file(src, dst)

        super(PythonPackageBuild, self).finalize_options()


class InstallCmd(install):
    def finalize_options(self) -> None:
        self.build_lib = os.path.relpath(os.path.join(BASE_DIR, "build"))
        return super(InstallCmd, self).finalize_options()


ext_modules = [
    Pybind11Extension(
        'npu_inductor_mlir._C', 
        ['npu_inductor_mlir/_C/extension.cpp'], 
        include_dirs=[
            'npu_inductor_mlir/_C/include',
            ],
        extra_link_args = ['-L' + os.path.join(asc_path, 'lib64'), '-lruntime'],
        extra_compile_args=["-std=c++17"],
        build_lib='dist',
    ),
]

setup(
    name='npu_inductor_mlir',
    version='0.0.1',
    author='Ascend',
    author_email='ascend@huawei.com',
    description='A project integrating PyTorch with MLIR for NPU acceleration.',
    ext_modules=ext_modules,
    cmdclass={
        'build_clib': CPPLibBuild,
        'build_ext': Build,
        'build_py': PythonPackageBuild,
        'install': InstallCmd,
        'clean': clean,
        'uninstall': CustomUninstallCommand,
    },
    packages=["npu_inductor_mlir"],
    package_data={
        'npu_inductor_mlir': [
            '*.so', 'lib/*.so'
        ],
    },
    include_package_data=True,
    exclude_package_data={
        'npu_inductor_mlir': ['*.cpp'],
    },
    install_requires =[
        'numpy < 2.0',
        'pybind11',
        'wheel',
        'setuptools',
        'cmake',
        'ninja',
        'packaging',
        'pyyaml',
    ],
    zip_safe=False
)