#! /usr/bin/env python3
import sys
import os
import shutil
import pathlib
import multiprocessing

from math import ceil

from setuptools import setup, Extension, find_packages
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext

#######################################################
# DEPENDENCIES
third_party_dependencies = ["numpy>=1.21.6", "Jinja2>=3.1.2"]
aidge_optional_dependencies = {"test": ["pytest"]}


class CMakeExtension(Extension):
    def __init__(self, name):
        super().__init__(name, sources=[])


class CMakeBuild(build_ext):
    def run(self):
        # This lists the number of processors available on the machine
        # The compilation will use half of them
        max_jobs = str(ceil(multiprocessing.cpu_count() / 2))

        cwd = pathlib.Path().absolute()

        build_temp = cwd / "build"
        if not build_temp.exists():
            build_temp.mkdir(parents=True, exist_ok=True)

        build_lib = pathlib.Path(self.build_lib)
        if not build_lib.exists():
            build_lib.mkdir(parents=True, exist_ok=True)

        os.chdir(str(build_temp))

        compile_type = (
            "Release"
            if "AIDGE_PYTHON_BUILD_TYPE" not in os.environ
            else os.environ["AIDGE_PYTHON_BUILD_TYPE"]
        )
        install_path = (
            os.environ["AIDGE_INSTALL"]
            if "AIDGE_INSTALL" in os.environ
            else os.path.join(sys.prefix, "lib", "libAidge")
        )
        self.spawn(
            [
                "cmake",
                str(cwd),
                "-DTEST=OFF",
                f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
                f"-DCMAKE_BUILD_TYPE={compile_type}",
                "-DPYBIND=ON",
                "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
                "-DCOVERAGE=OFF",
            ]
        )

        if not self.dry_run:
            self.spawn(
                ["cmake", "--build", ".", "--config", compile_type, "-j", max_jobs]
            )
            self.spawn(["cmake", "--install", ".", "--config", compile_type])
        os.chdir(str(cwd))

        aidge_package = build_lib / "aidge_core"
        print(f"aidge_package path = {aidge_package.absolute()}")

        # Get "aidge core" package
        print(f"Build directory = {build_temp.absolute()}")
        # Copy all shared object files from build_temp/lib to aidge_package
        for root, _, files in os.walk(build_temp.absolute()):
            for file in files:
                if (file.endswith(".so") or file.endswith(".pyd")) and (
                    root != str(aidge_package.absolute())
                ):
                    currentFile = os.path.join(root, file)
                    print(f"currentFile = {currentFile}")
                    shutil.copy(currentFile, str(aidge_package.absolute()))


def clean_eggs():
    eggs_dir = ".eggs"
    if os.path.isdir(eggs_dir):
        shutil.rmtree(eggs_dir)
        print(f"Removed {eggs_dir} directory")


if __name__ == "__main__":
    setup(
        name="aidge_core",
        description="Core algorithms for operators and graph of the AIDGE framework",
        python_requires=">=3.7",
        setup_requires=[
            "setuptools>=64",
            "setuptools_scm[toml]==7.1.0",
            "cmake>=3.27.9",
        ],
        install_require=third_party_dependencies,
        extras_require=aidge_optional_dependencies,
        readme="README.md",
        license="LICENSE",
        classifiers=[
            "Development Status :: 2 - Pre-Alpha",
            "Programming Language :: Python :: 3",
        ],
        use_scm_version=True,
        packages=find_packages(where="."),
        include_package_data=True,
        ext_modules=[CMakeExtension("aidge_core")],
        cmdclass={
            "build_ext": CMakeBuild,
        },
        zip_safe=False,
    )
    clean_eggs()