import os
import re
import sys
import sysconfig
import platform
import subprocess
import multiprocessing

from distutils.version import LooseVersion
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.install_lib import install_lib

PACKAGE_NAME = "baggianalysis"


class CMakeExtension(Extension):

    def __init__(self, name, sourcedir=''):
        Extension.__init__(self, name, sources=[])
        self.sourcedir = os.path.abspath(sourcedir)
        
        
class InstallCMakeLibs(install_lib):
    """
    Get the libraries from the parent distribution, use those as the outfiles

    Skip building anything; everything is already built, forward libraries to
    the installation step
    """

    def run(self):
        """
        Copy libraries from the bin directory and place them as appropriate
        """

        self.announce("Moving library files", level=3)

        # We have already built the libraries in the previous build_ext step
        self.skip_build = True

        # Folder where the `baggianalysis` package has been placed by cmake. It is used by self.install
        self.build_dir = self.distribution.lib_dir
        self.outfiles = self.install()
        
        # I have copied this bit from the parent class
        if self.outfiles is not None:
            # always compile, in case we have any extension stubs to deal with
            self.byte_compile(self.outfiles)
            
    def get_outputs(self):
        """
        Overrides the parent class' method. Returns a list of the files copied over by the `run` method 
        """
        return self.outfiles


class CMakeBuild(build_ext):

    def run(self):
        try:
            out = subprocess.check_output(['cmake', '--version'])
        except OSError:
            raise RuntimeError("CMake must be installed to build the following extensions: " + ", ".join(e.name for e in self.extensions))

        self.cmake_version = LooseVersion(re.search(r'version\s*([\d.]+)', out.decode()).group(1))

        if platform.system() == "Windows":
            if self.cmake_version < '3.1.0':
                raise RuntimeError("CMake >= 3.1.0 is required on Windows")

        for ext in self.extensions:
            self.build_extension(ext)
            
    def build_extension(self, ext):
        
        self.announce("Preparing the build environment", level=3)
        
        cmake_args = []

        cfg = 'Debug' if self.debug else 'Release'
        build_args = ['--config', cfg]
        native_generator_args = ['--']

        if platform.system() == "Windows":
            if sys.maxsize > 2 ** 32:
                cmake_args += ['-A', 'x64']
            native_generator_args += ['/m']
        else:
            cmake_args += ['-DCMAKE_BUILD_TYPE=' + cfg]
                
        cmake_args += ['-DPython=On']
        
        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level across all generators.
        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
            
            try:
                cpu_cores = int(os.getenv('SLURM_NTASKS'))
            except:
                cpu_cores = int(multiprocessing.cpu_count() / 2)
                
            if self.cmake_version < "3.14.0":
                native_generator_args += ["-j{}".format(cpu_cores)]
            else:
                build_args += ["-j {}".format(cpu_cores)]
                
        build_args += native_generator_args

        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)
            
        self.distribution.lib_dir = os.path.join(self.build_temp, "lib/python")
            
        self.announce("Configuring cmake project", level=3)
        subprocess.check_call(['cmake', ext.sourcedir] + cmake_args, cwd=self.build_temp)
        
        self.announce("Building the library", level=3)
        subprocess.check_call(['cmake', '--build', '.'] + build_args, cwd=self.build_temp)

        self.announce("Compilation done", level=3)
        

setup(
    name = PACKAGE_NAME,
    use_scm_version = {
        "fallback_version": "0.0.2",
        },
    packages = find_packages(),
    setup_requires = [
        'setuptools-scm',
        'numpy'
        ],
    author = 'Lorenzo Rovigatti',
    author_email = 'lorenzo.rovigatti@uniroma1.it',
    url = 'https://github.com/lorenzo-rovigatti/baggianalysis',
    description = 'A C++/Python library to facilitate the analysis of molecular simulations',
    long_description = open("./README.md", 'r').read(),
    long_description_content_type = "text/markdown",
    license = 'GNU GPL 3.0',
    ext_modules = [CMakeExtension('baggianalysis')],
    # add custom build_ext command
    cmdclass = {
        'build_ext': CMakeBuild,
        'install_lib': InstallCMakeLibs,
        },
    zip_safe = False,
    classifiers = [
        "Development Status :: 4 - Beta",
        "Environment :: Console",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
        "Operating System :: OS Independent",
        "Programming Language :: C++",
        "Topic :: Scientific/Engineering :: Physics"
        ],
    entry_points = {
        'console_scripts': 
            [
                'oxDNA_LAMMPS=baggianalysis.converters:oxDNA_LAMMPS_command_line', 
                'oxDNA_cogli=baggianalysis.converters:oxDNA_cogli_command_line', 
                'LAMMPS_cogli=baggianalysis.converters:LAMMPS_cogli_command_line'
            ],
    },
)
