"""
 Copyright (C) 2018-2021 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 http://www.apache.org/licenses/LICENSE-2.0
 This conversation was marked as resolved by dkurt
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import os.path
import sys
import errno
import subprocess
from pathlib import Path
from shutil import copyfile
from distutils.command.install import install
from distutils.command.build import build
from distutils.errors import DistutilsSetupError
from distutils.file_util import copy_file
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_clib import build_clib
from decouple import config

WHEEL_LIBS_INSTALL_DIR = os.path.join('openvino', 'libs')
WHEEL_LIBS_PACKAGE = 'openvino.libs'
PYTHON_VERSION = f"python{sys.version_info.major}.{sys.version_info.minor}"

# The following variables can be defined in environment or .env file
CMAKE_BUILD_DIR = config('CMAKE_BUILD_DIR', ".")
CORE_LIBS_DIR = config('CORE_LIBS_DIR', '')
PLUGINS_LIBS_DIR = config('PLUGINS_LIBS_DIR', '')
NGRAPH_LIBS_DIR = config('NGRAPH_LIBS_DIR', '')
TBB_LIBS_DIR = config('TBB_LIBS_DIR', '')
PY_PACKAGES_DIR = config('PY_PACKAGES_DIR', '')
LIBS_RPATH = "$ORIGIN" if sys.platform == "linux" else "@loader_path"

LIB_INSTALL_CFG = {
    "ie_libs": {
        'name': 'core',
        'prefix': 'libs.core',
        'install_dir': CORE_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "hetero_plugin": {
        'name': 'hetero',
        'prefix': 'libs.plugins',
        'install_dir': PLUGINS_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "gpu_plugin": {
        'name': 'gpu',
        'prefix': 'libs.plugins',
        'install_dir': PLUGINS_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "cpu_plugin": {
        'name': 'cpu',
        'prefix': 'libs.plugins',
        'install_dir': PLUGINS_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "multi_plugin": {
        'name': 'multi',
        'prefix': 'libs.plugins',
        'install_dir': PLUGINS_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "myriad_plugin": {
        'name': 'myriad',
        'prefix': 'libs.plugins',
        'install_dir': PLUGINS_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "ngraph_libs": {
        'name': 'ngraph',
        'prefix': 'libs.ngraph',
        'install_dir': NGRAPH_LIBS_DIR,
        'rpath': LIBS_RPATH,
    },
    "tbb_libs": {
        'name': 'tbb',
        'prefix': 'libs.tbb',
        'install_dir': TBB_LIBS_DIR
    },
}

PY_INSTALL_CFG = {
    "ie_py": {
        'name': PYTHON_VERSION,
        'prefix': 'site-packages',
        'install_dir': PY_PACKAGES_DIR
    },
    "ngraph_py": {
        'name': f"pyngraph_{PYTHON_VERSION}",
        'prefix': 'site-packages',
        'install_dir': PY_PACKAGES_DIR
    },
}


class PrebuiltExtension(Extension):
    """Initialize Extension"""
    def __init__(self, name, sources, *args, **kwargs):
        if len(sources) != 1:
            nln = '\n'
            raise DistutilsSetupError(
                f"PrebuiltExtension can accept only one source, but got: {nln}{nln.join(sources)}"
            )
        super().__init__(name, sources, *args, **kwargs)


class CustomBuild(build):
    """Custom implementation of build_clib"""
    def run(self):
        self.run_command('build_clib')
        build.run(self)


class CustomInstall(install):
    """Enable build_clib during the installation"""
    def run(self):
        self.run_command('build_clib')
        install.run(self)


class PrepareLibs(build_clib):
    """Prepare prebuilt libraries"""
    def run(self):
        self.configure(LIB_INSTALL_CFG)
        self.configure(PY_INSTALL_CFG)
        self.generate_package(get_dir_list(LIB_INSTALL_CFG))

    def configure(self, install_cfg):
        """Collect prebuilt libraries. Install them to the temp directories, set rpath."""
        for comp, comp_data in install_cfg.items():
            install_prefix = comp_data.get('prefix')
            install_dir = comp_data.get('install_dir')
            if install_dir and not os.path.isabs(install_dir):
                install_dir = os.path.join(install_prefix, install_dir)
                self.announce(f"Installing {comp}", level=3)
                self.spawn(["cmake",
                            "--install", CMAKE_BUILD_DIR,
                            "--prefix", install_prefix,
                            "--component", comp_data.get('name')])
            # set rpath if applicable
            if sys.platform != "win32" and comp_data.get('rpath'):
                lib_pattern = "*.so" if sys.platform == "linux" else "*.dylib"
                for path in Path(install_dir).glob(lib_pattern):
                    set_rpath(comp_data['rpath'], path)

    def generate_package(self, src_dirs):
        """
        Collect package data files from preinstalled dirs and
        put all runtime libraries to the subpackage
        """
        # additional blacklist filter, just to fix cmake install issues
        blacklist = ['.lib', '.pdb', '_debug.dll', '_debug.dylib']
        package_dir = os.path.join(get_package_dir(PY_INSTALL_CFG), WHEEL_LIBS_INSTALL_DIR)
        for src_dir in src_dirs:
            local_base_dir = Path(src_dir)
            for file_path in local_base_dir.rglob('*'):
                file_name = os.path.basename(file_path)
                if file_path.is_file() and not any(file_name.endswith(ext) for ext in blacklist):
                    dst_file = os.path.join(package_dir, os.path.relpath(file_path, local_base_dir))
                    os.makedirs(os.path.dirname(dst_file), exist_ok=True)
                    copyfile(file_path, dst_file)

        if Path(package_dir).exists():
            self.announce(f"Adding {WHEEL_LIBS_PACKAGE} package", level=3)
            packages.append(WHEEL_LIBS_PACKAGE)
            package_data.update({WHEEL_LIBS_PACKAGE: ['*']})


class CopyExt(build_ext):
    """Copy extension files to the build directory"""
    def run(self):
        for extension in self.extensions:
            if not isinstance(extension, PrebuiltExtension):
                raise DistutilsSetupError(
                    f"copy_ext can accept PrebuiltExtension only, but got {extension.name}")
            src = extension.sources[0]
            dst = self.get_ext_fullpath(extension.name)
            os.makedirs(os.path.dirname(dst), exist_ok=True)
            # setting relative path to find dlls
            if sys.platform != "win32":
                rpath = os.path.relpath(get_package_dir(PY_INSTALL_CFG), os.path.dirname(src))
                if sys.platform == "linux":
                    rpath = os.path.join("$ORIGIN", rpath, WHEEL_LIBS_INSTALL_DIR)
                elif sys.platform == "darwin":
                    rpath = os.path.join("@loader_path", rpath, WHEEL_LIBS_INSTALL_DIR)
                set_rpath(rpath, src)

            copy_file(src, dst, verbose=self.verbose, dry_run=self.dry_run)


def is_tool(name):
    """Check if the command-line tool is available"""
    try:
        devnull = subprocess.DEVNULL
        subprocess.Popen([name], stdout=devnull, stderr=devnull).communicate()
    except OSError as error:
        if error.errno == errno.ENOENT:
            return False
    return True


def set_rpath(rpath, executable):
    """Setting rpath for linux and macOS libraries"""
    print(f"Setting rpath {rpath} for {executable}")
    cmd = ""
    rpath_tool = ""
    if sys.platform == "linux":
        rpath_tool = "patchelf"
        cmd = [rpath_tool, "--set-rpath", rpath,  executable]
    elif sys.platform == "darwin":
        rpath_tool = "install_name_tool"
        cmd = [rpath_tool, "-add_rpath", rpath,  executable]
    else:
        sys.exit(f"Unsupported platform: {sys.platform}")

    if is_tool(rpath_tool):
        ret_info = subprocess.run(cmd, check=True)
        if ret_info.returncode != 0:
            sys.exit(f"Could not set rpath: {rpath} for {executable}")
    else:
        sys.exit(f"Could not found {rpath_tool} on the system, "
                 f"please make sure that this tool is installed")


def find_prebuilt_extensions(search_dirs):
    """collect prebuilt python extensions"""
    extensions = []
    ext_pattern = ""
    if sys.platform == "linux":
        ext_pattern = "**/*.so"
    elif sys.platform == "win32":
        ext_pattern = "**/*.pyd"
    elif sys.platform == "darwin":
        ext_pattern = "**/*.so"
    for base_dir in search_dirs:
        for path in Path(base_dir).glob(ext_pattern):
            relpath = path.relative_to(base_dir)
            if relpath.parent != ".":
                package_names = str(relpath.parent).split(os.path.sep)
            else:
                package_names = []
            package_names.append(path.name.split(".", 1)[0])
            name = ".".join(package_names)
            extensions.append(PrebuiltExtension(name, sources=[str(path)]))
    return extensions


def get_description(desc_file_path):
    """read description from README.md"""
    with open(desc_file_path, "r", encoding="utf-8") as fstream:
        description = fstream.read()
    return description


def get_dependencies(requirements_file_path):
    """read dependencies from requirements.txt"""
    with open(requirements_file_path, "r", encoding="utf-8") as fstream:
        dependencies = fstream.read()
    return dependencies


def get_dir_list(install_cfg):
    """collect all available directories with libs or python packages"""
    dirs = []
    for comp_info in install_cfg.values():
        cfg_prefix = comp_info.get('prefix')
        cfg_dir = comp_info.get('install_dir')
        if cfg_dir:
            if not os.path.isabs(cfg_dir):
                cfg_dir = os.path.join(cfg_prefix, cfg_dir)
            if cfg_dir not in dirs:
                dirs.append(cfg_dir)
    return dirs


def get_package_dir(install_cfg):
    """
    Get python package path based on config
    All the packages should be located in one directory
    """
    py_package_path = ""
    dirs = get_dir_list(install_cfg)
    if len(dirs) != 0:
        # setup.py support only one package directory, all modules should be located there
        py_package_path = dirs[0]
    return py_package_path


platforms = ["linux", "win32", "darwin"]
if not any(pl in sys.platform for pl in platforms):
    sys.exit("Unsupported platform: {}, expected: {}".format(sys.platform, "linux, win32, darwin"))

# copy license file into the build directory
package_license = config('WHEEL_LICENSE', '')
if os.path.exists(package_license):
    copyfile(package_license, "LICENSE")

packages = find_packages(','.join(get_dir_list(PY_INSTALL_CFG)))
package_data = {}

setup(
    version=config('WHEEL_VERSION', '0.0.0'),
    author_email=config('WHEEL_AUTHOR_EMAIL', 'openvino_pushbot@intel.com'),
    name=config('WHEEL_PACKAGE_NAME', 'openvino'),
    license=config('WHEEL_LICENCE_TYPE', 'OSI Approved :: Apache Software License'),
    author=config('WHEEL_AUTHOR', 'Intel Corporation'),
    description=config('WHEEL_DESC', 'Inference Engine Python* API'),
    install_requires=get_dependencies(config('WHEEL_REQUIREMENTS', "requirements.txt")),
    long_description=get_description(config('WHEEL_OVERVIEW', 'pypi_overview.md')),
    long_description_content_type="text/markdown",
    download_url=config('WHEEL_DOWNLOAD_URL', 'https://github.com/openvinotoolkit/openvino/tags'),
    url=config('WHEEL_URL', 'https://docs.openvinotoolkit.org/latest/index.html'),
    cmdclass={
        "build": CustomBuild,
        "install": CustomInstall,
        "build_clib": PrepareLibs,
        "build_ext": CopyExt,
    },
    ext_modules=find_prebuilt_extensions(get_dir_list(PY_INSTALL_CFG)),
    packages=packages,
    package_dir={'': get_package_dir(PY_INSTALL_CFG)},
    package_data=package_data,
    zip_safe=False,
)
