import importlib
import os
import sys
from distutils.command.clean import clean
from pathlib import Path
from typing import Union, List

import setuptools
from setuptools.extension import Extension


root_path = Path(__file__).resolve().parent


version = "unknown"
version_file = root_path / "version.txt"
if version_file.exists():
    with open(root_path / "version.txt") as handle:
        version_content = handle.read().strip()
        if version_content:
            version = version_content
print("Bitorch Engine: version", version)


ALL_EXTENSIONS = [
    "bitorch_engine/layers/qconv/binary/cpp",
    "bitorch_engine/layers/qconv/binary/cuda",
    "bitorch_engine/layers/qconv/nbit/cutlass",
    "bitorch_engine/layers/qlinear/binary/cpp",
    "bitorch_engine/layers/qlinear/binary/cutlass",
    "bitorch_engine/layers/qlinear/binary/cuda",
    "bitorch_engine/layers/qlinear/nbit/cutlass",
    "bitorch_engine/layers/qembedding/binary/cuda",
    "bitorch_engine/functions/cuda",
]


def get_requirements(file_path: Union[Path, str]):
    requires = []
    for requirement in (root_path / file_path).open().readlines():
        requires.append(requirement)
    return requires


dependencies = get_requirements('requirements.txt')


class CustomClean(clean):
    def run(self):
        # clean all by default
        self.all = True
        super().run()
        import bitorch_engine.extensions as ext
        for file in Path(ext.__file__).parent.glob("*.pyd" if sys.platform == "win32" else "*.so"):
            if self.dry_run:
                print(f"would remove '{file}'")
            else:
                print(f"removing '{file}'")
                file.unlink()


def get_ext_modules() -> List[Extension]:
    dirs_to_process = ALL_EXTENSIONS
    if "BIE_BUILD_ONLY" in os.environ:
        dirs_to_process = os.environ.get("BIE_BUILD_ONLY").split(" ")
        print(f"Extensions to build from environment variable 'BIE_BUILD_ONLY': {dirs_to_process}")
        for extension_dir in dirs_to_process:
            if extension_dir not in ALL_EXTENSIONS:
                sep = "\n    - "
                print(f"The extension directory '{extension_dir}' is not valid.\n" +
                      "Must be one of:{}{}".format(sep, sep.join(ALL_EXTENSIONS)))
                sys.exit(1)
    extension_list = []
    cutlass_installed = is_cutlass_available()
    for extension_dir in dirs_to_process:
        # we can print debug information here, it is only shown in case of errors or with a verbose setting
        print(f"Loading extension for build from {extension_dir}")
        sys.path.insert(0, extension_dir)
        module = importlib.import_module(extension_dir.replace("/", ".") + ".extension")
        if hasattr(module, "CUDA_REQUIRED") and module.CUDA_REQUIRED and not torch.cuda.is_available():
            print(f"Skipping {module}, since it seems as if CUDA is not available.")
            continue
        if hasattr(module, "CUTLASS_REQUIRED") and module.CUTLASS_REQUIRED and not cutlass_installed:
            print(f"Skipping {module}, since it seems as if CUTLASS is not available.")
            continue
        extension_list.append(module.get_ext(Path(extension_dir)))
        sys.path.pop(0)
    return extension_list


cmdclass = {'clean': CustomClean}
ext_modules = []
if "develop" not in sys.argv:
    print("Bitorch Engine: Not loading extension build commands.")
else:
    print("Bitorch Engine: Loading extension build commands.")
    # we have installed requirements already and (might) need to build extensions
    import torch
    from torch.utils.cpp_extension import BuildExtension, IS_LINUX
    from bitorch_engine.utils.cutlass_path import is_cutlass_available
    from bitorch_engine.utils.cuda_extension import gcc_version

    class CustomBuildExtension(BuildExtension):
        def run(self):
            if IS_LINUX:
                major, minor, patch = gcc_version()
                print("GCC Version:", f"{major}.{minor}.{patch}")
            print("Python Path:", sys.path)
            print("Environment:\n", "\n    ".join(f"{k}='{v}'" for k, v in os.environ.items()))
            super().run()

    cmdclass['build_ext'] = CustomBuildExtension
    ext_modules = get_ext_modules()


with open("README.md", "r", encoding="utf-8") as handle:
    readme_content = handle.read()


setuptools.setup(
    name="bitorch_engine",
    url="https://github.com/hpi-xnor/bitorch-inference-engine",
    version=version,
    author="Hasso Plattner Institute",
    author_email="fb10-xnor@hpi.de",
    description="A package for building and training quantized and binary neural networks with Pytorch",
    long_description=readme_content,
    long_description_content_type="text/markdown",
    packages=setuptools.find_packages(exclude='tests'),
    install_requires=dependencies,
    extras_require={
        "dev": get_requirements("requirements-dev.txt"),
    },
    classifiers=[
        "Development Status :: 3 - Alpha",
        "Intended Audience :: Developers",
        "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3 :: Only",
        "Programming Language :: Python :: 3.7",
    ],
    python_requires='>=3.7',
    data_files=[
        ('.', [
            'version.txt',
            'requirements.txt',
            'requirements-dev.txt',
        ]),
    ],
    ext_modules=ext_modules,
    cmdclass=cmdclass,  # type: ignore
)
