# Copyright © 2019 Arm Ltd. All rights reserved.
# SPDX-License-Identifier: MIT
import logging
import os
import sys
from functools import lru_cache
from pathlib import Path
from itertools import chain

from setuptools import setup
from distutils.core import Extension
from setuptools.command.build_py import build_py
from setuptools.command.build_ext import build_ext

logger = logging.Logger(__name__)

__version__ = None
__arm_ml_version__ = None


def check_armnn_version(*args):
    pass


exec(open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src', 'pyarmnn', '_version.py')).read())


class ExtensionPriorityBuilder(build_py):
    """
    Runs extension builder before other stages. Otherwise generated files are not included to the distribution.
    """

    def run(self):
        self.run_command('build_ext')
        return super().run()


class ArmnnVersionCheckerExtBuilder(build_ext):

    def __init__(self, dist):
        super().__init__(dist)
        self.failed_ext = []

    def build_extension(self, ext):
        try:
            super().build_extension(ext)
        except Exception as err:
            self.failed_ext.append(ext)
            logger.warning('Failed to build extension %s. \n %s', ext.name, str(err))

        if ext.name == 'pyarmnn._generated._pyarmnn_version':
            sys.path.append(os.path.abspath(os.path.join(self.build_lib, str(Path(ext._file_name).parent))))
            from _pyarmnn_version import GetVersion
            check_armnn_version(GetVersion(), __arm_ml_version__)

    def copy_extensions_to_source(self):

        for ext in self.failed_ext:
            self.extensions.remove(ext)
        super().copy_extensions_to_source()


def linux_gcc_lib_search():
    """
    Calls the `gcc` to get linker default system paths.
    Returns:
        list of paths
    """
    cmd = 'gcc --print-search-dirs | grep libraries'
    cmd_res = os.popen(cmd).read()
    cmd_res = cmd_res.split('=')
    if len(cmd_res) > 1:
        return tuple(cmd_res[1].split(':'))
    return None


def find_includes(armnn_include_env: str = 'ARMNN_INCLUDE'):
    armnn_include_path = os.getenv(armnn_include_env, '')
    return [armnn_include_path] if armnn_include_path else ['/usr/local/include', '/usr/include']


@lru_cache(maxsize=1)
def find_armnn(lib_name: str,
               optional: bool = False,
               armnn_libs_env: str = 'ARMNN_LIB',
               default_lib_search: tuple = linux_gcc_lib_search()):
    """
    Searches for ArmNN installation on the local machine.

    Args:
        lib_name: lib name to find
        optional: Do not fail if optional. Default is False - fail if library was not found.
        armnn_include_env: custom environment variable pointing to ArmNN headers, default is 'ARMNN_INCLUDE'
        armnn_libs_env: custom environment variable pointing to ArmNN libraries location, default is 'ARMNN_LIBS'
        default_lib_search: list of paths to search for ArmNN if not found within path provided by 'ARMNN_LIBS'
                            env variable

    Returns:
        tuple containing name of the armnn libs, paths to the libs
    """

    armnn_lib_path = os.getenv(armnn_libs_env, "")

    lib_search = [armnn_lib_path] if armnn_lib_path else default_lib_search

    armnn_libs = dict(map(lambda path: (':{}'.format(path.name), path),
                          chain.from_iterable(map(lambda lib_path: Path(lib_path).glob(lib_name),
                                                  lib_search))))
    if not optional and len(armnn_libs) == 0:
        raise RuntimeError("""ArmNN library {} was not found in {}. Please install ArmNN to one of the standard
                           locations or set correct ARMNN_INCLUDE and ARMNN_LIB env variables.""".format(lib_name,
                                                                                                         lib_search))

    # gives back tuple of names of the libs, set of unique libs locations and includes.
    return list(armnn_libs.keys()), list(set(
        map(lambda path: str(path.absolute().parent), armnn_libs.values())))


class LazyArmnnFinderExtension(Extension):
    """
    Derived from `Extension` this class adds ArmNN libraries search on the user's machine.
    SWIG options and compilation flags are updated with relevant ArmNN libraries files locations (-L) and headers (-I).

    Search for ArmNN is executed only when attributes include_dirs, library_dirs, runtime_library_dirs, libraries or
    swig_opts are queried.

    """

    def __init__(self, name, sources, armnn_libs, include_dirs=None, define_macros=None, undef_macros=None,
                 library_dirs=None,
                 libraries=None, runtime_library_dirs=None, extra_objects=None, extra_compile_args=None,
                 extra_link_args=None, export_symbols=None, language=None, optional=None, **kw):
        self._include_dirs = None
        self._library_dirs = None
        self._runtime_library_dirs = None
        self._armnn_libs = armnn_libs
        # self.__swig_opts = None
        super().__init__(name, sources, include_dirs, define_macros, undef_macros, library_dirs, libraries,
                         runtime_library_dirs, extra_objects, extra_compile_args, extra_link_args, export_symbols,
                         language, optional, **kw)

    @property
    def include_dirs(self):
        return self._include_dirs + find_includes()

    @include_dirs.setter
    def include_dirs(self, include_dirs):
        self._include_dirs = include_dirs

    @property
    def library_dirs(self):
        library_dirs = self._library_dirs
        for lib in self._armnn_libs:
            _, lib_path = find_armnn(lib)
            library_dirs = library_dirs + lib_path

        return library_dirs

    @library_dirs.setter
    def library_dirs(self, library_dirs):
        self._library_dirs = library_dirs

    @property
    def runtime_library_dirs(self):
        library_dirs = self._runtime_library_dirs
        for lib in self._armnn_libs:
            _, lib_path = find_armnn(lib)
            library_dirs = library_dirs + lib_path

        return library_dirs

    @runtime_library_dirs.setter
    def runtime_library_dirs(self, runtime_library_dirs):
        self._runtime_library_dirs = runtime_library_dirs

    @property
    def libraries(self):
        libraries = self._libraries
        for lib in self._armnn_libs:
            lib_names, _ = find_armnn(lib)
            libraries = libraries + lib_names

        return libraries

    @libraries.setter
    def libraries(self, libraries):
        self._libraries = libraries

    def __eq__(self, other):
        return self.__class__ == other.__class__ and self.name == other.name

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return self.name.__hash__()

if __name__ == '__main__':
    # mandatory extensions
    pyarmnn_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn',
                                              sources=['src/pyarmnn/_generated/armnn_wrap.cpp'],
                                              extra_compile_args=['-std=c++14'],
                                              language='c++',
                                              armnn_libs=['libarmnn.so']
                                              )
    pyarmnn_v_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn_version',
                                                sources=['src/pyarmnn/_generated/armnn_version_wrap.cpp'],
                                                extra_compile_args=['-std=c++14'],
                                                language='c++',
                                                armnn_libs=['libarmnn.so']
                                                )
    extensions_to_build = [pyarmnn_v_module, pyarmnn_module]


    # optional extensions
    def add_parsers_ext(name: str, ext_list: list):
        pyarmnn_optional_module = LazyArmnnFinderExtension('pyarmnn._generated._pyarmnn_{}'.format(name.lower()),
                                                           sources=['src/pyarmnn/_generated/armnn_{}_wrap.cpp'.format(
                                                               name.lower())],
                                                           extra_compile_args=['-std=c++14'],
                                                           language='c++',
                                                           armnn_libs=['libarmnn.so', 'libarmnn{}.so'.format(name)]
                                                           )
        ext_list.append(pyarmnn_optional_module)


    add_parsers_ext('CaffeParser', extensions_to_build)
    add_parsers_ext('OnnxParser', extensions_to_build)
    add_parsers_ext('TfParser', extensions_to_build)
    add_parsers_ext('TfLiteParser', extensions_to_build)

    setup(
        name='pyarmnn',
        version=__version__,
        author='Arm ltd',
        description='Arm NN python wrapper',
        url='https://www.arm.com',
        license='MIT',
        package_dir={'': 'src'},
        packages=[
            'pyarmnn',
            'pyarmnn._generated',
            'pyarmnn._quantization',
            'pyarmnn._tensor',
            'pyarmnn._utilities'
        ],
        data_files=[('', ['licences.txt'])],
        python_requires='>=3.5',
        install_requires=['numpy'],
        cmdclass={'build_py': ExtensionPriorityBuilder, 'build_ext': ArmnnVersionCheckerExtBuilder},
        ext_modules=extensions_to_build
    )