from pathlib import Path
import os
import stat
import sysconfig
import shutil
import subprocess
import glob
import torch
import torch_npu
from setuptools import setup
from setuptools import Extension
from setuptools.command.build_clib import build_clib
from distutils.command.build_py import build_py
from setuptools.command.build_ext import build_ext
from setuptools.command.install import install
from setuptools.command.build_clib import build_clib

VERSION = '0.1.0'
BASE_DIR = os.path.dirname(os.path.realpath(__file__))

def _get_ascend_path() -> str:
    path = os.getenv("ASCEND_HOME_PATH", "")
    if path == "":
        raise Exception("ASCEND_HOME_PATH is not set, source <ascend-toolkit>/set_env.sh first")
    return Path(path)


def _build_npu_common():
    build_dir = os.path.join(BASE_DIR, "build")
    lib_dir = os.path.join(build_dir, "npu_inductor/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/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')}"] # hacl_rt.h
    asc_path = _get_ascend_path()
    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 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", '*.py'),
            recursive=True) + glob.glob(
            os.path.join(BASE_DIR, "npu_inductor", '**/*.py'),
            recursive=True) + glob.glob(
            os.path.join(BASE_DIR, "npu_inductor", '*.h'),
            recursive=True)
        for src in src_files:
            dst = os.path.join(
                os.path.join(BASE_DIR, "build/npu_inductor"),
                os.path.relpath(src, os.path.join(BASE_DIR, "npu_inductor")))
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            self.copy_file(src, dst)
        super(PythonPackageBuild, self).finalize_options()


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


class Build(build_ext):
    def run(self):
        self.run_command('build_clib')
        super(Build, self).run()


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()


setup(
    name='npu_inductor',
    version=VERSION,
    description='NPU bridge for PyTorch Inductor',
    packages=["npu_inductor"],
    package_dir={'': os.path.relpath(os.path.join(BASE_DIR, "build"))},
    ext_modules=[Extension("dummy.extension", sources=[])],
    package_data={
        'npu_inductor': [
            'lib/*.so',
        ],
    },
    cmdclass={
        'build_clib': CPPLibBuild,
        'build_ext': Build,
        'build_py': PythonPackageBuild,
        'install': InstallCmd
    })