import os
import sys
import subprocess
import multiprocessing
import shutil

from pathlib import Path

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


__version__ = "0.0.4"
HERE = Path(__file__).resolve().parent


# A CMakeExtension needs a sourcedir instead of a file list.
# The name must be the _single_ output extension from the CMake build.
# If you need multiple extensions, see scikit-build.
class CMakeExtension(Extension):
    def __init__(self, name, sourcedir=""):
        Extension.__init__(self, name, sources=[])
        self.sourcedir = Path(os.path.abspath(sourcedir))


def get_num_jobs():
    return int(os.environ.get("TYR_JOBS", multiprocessing.cpu_count()))

def get_header_instantiation():
    return os.environ.get("TYR_HEADER_INSTANTIATION", "Off")


class CMakeBuild(build_ext):
    def build_extension(self, ext):
        # Must be in this form due to bug in .resolve() only fixed in Python 3.10+
        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)
        output_directory = ext_fullpath.parent.resolve()
        temp_directory = Path.cwd() / self.build_temp

        print("ext_fullpath", ext_fullpath)
        print("output_directory", output_directory)
        print("temp_directory", temp_directory)

        build_type = "Debug" if os.environ.get('TYR_DEBUG_BUILD') else "Release"
        print("tyr build type:", build_type)

        # Create the temporary build directory, if it does not already exist
        os.makedirs(temp_directory, exist_ok=True)

        #######################################################################
        # Build dependencies
        #######################################################################

        cmake_args = [
            f"-DCMAKE_BUILD_TYPE={build_type}",
            f"-DCMAKE_INSTALL_PREFIX={str(temp_directory / 'dependencies' / 'installs')}",
            f"-DCMAKE_PREFIX_PATH={str(temp_directory / 'dependencies' / 'installs')}",
            f"-DPython_EXECUTABLE={sys.executable}"
        ]

        subprocess.run(
            ["cmake", "-S", f"{str(ext.sourcedir / 'dependencies')}", "-B", f"{str(temp_directory / 'dependencies' / 'build')}"] + cmake_args, cwd=str(temp_directory), check=True
        )

        subprocess.run(
            ["cmake", "--build", f"{str(temp_directory / 'dependencies' / 'build')}", f"-j{get_num_jobs()}"]
        )

        subprocess.run(
            ["cmake", "--install", f"{str(temp_directory / 'dependencies' / 'build')}"]
        )

        shutil.rmtree(f"{str(temp_directory / 'dependencies' / 'build')}")

        #######################################################################
        # Build tyr
        #######################################################################

        cmake_args = [
            f"-DTYR_VERSION_INFO={__version__}",
            f"-DCMAKE_BUILD_TYPE={build_type}",  # not used on MSVC, but no harm
            f"-DCMAKE_PREFIX_PATH={str(temp_directory / 'dependencies' / 'installs')}",
            f"-DPython_EXECUTABLE={sys.executable}",
            "-DBUILD_PYTYR=ON",
            "-DBUILD_TESTS=OFF",
            "-DBUILD_EXECUTABLES=OFF",
            "-DBUILD_PROFILING=OFF",
            "-DTYR_USE_LLD=OFF",  # lld seems to prune code
            f"-DTYR_HEADER_INSTANTIATION={get_header_instantiation()}"
        ]

        subprocess.run(
            ["cmake", "-S", ext.sourcedir, "-B", f"{str(temp_directory / 'build')}"] + cmake_args, cwd=str(temp_directory), check=True
        )

        subprocess.run(
            ["cmake", "--build", f"{str(temp_directory / 'build')}", f"-j{get_num_jobs()}"], cwd=str(temp_directory), check=True
        )

        subprocess.run(
            ["cmake", "--install", f"{str(temp_directory / 'build')}", "--prefix", f"{str(output_directory)}"], check=True
        )

        # Remove unwanted directories.
        unwanted_dirs = ["include", "lib", "lib64"]
        for unwanted_dir in unwanted_dirs:
            dir_path = output_directory / unwanted_dir
            if os.path.exists(dir_path):
                shutil.rmtree(dir_path)
                print(f"Removed {dir_path} from the wheel.")


# The information here can also be placed in setup.cfg - better separation of
# logic and declaration, and simpler if you include description/version in a file.
setup(
    name="pytyr",
    version=__version__,
    author="Dominik Drexler",
    author_email="dominik.drexler@liu.se",
    url="https://github.com/drexlerd/tyr",
    description="Tyr datalog/planning library",
    long_description="",
    install_requires=[],
    packages=find_packages(where="python/src"),
    package_dir={"": "python/src"},
    ext_modules=[CMakeExtension("pytyr")],
    cmdclass={"build_ext": CMakeBuild},
    zip_safe=False,
    extras_require={
        "test": [
            "pytest",
        ],
    }
)
