# Copyright (c) 2016-2023 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
#
# 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 multiprocessing.pool
import ctypes
import atexit
import sys
import os

import platform

if "Windows" in platform.system():
    import site
    path_to_env = site.getsitepackages()[0]
    path_to_libs = os.path.join(path_to_env, "Library", "bin")
    if sys.version_info.minor >= 8:
        os.add_dll_directory(path_to_libs)
    os.environ['PATH'] += os.pathsep + path_to_libs


from .api import  *
from .api import __all__ as api__all
from .pool import *
from .pool import __all__ as pool__all

__all__ = ["Monkey", "is_active"] + api__all + pool__all

__doc__ = """
Python API for Intel(R) oneAPI Threading Building Blocks (oneTBB)
extended with standard Python's pools implementation and monkey-patching.

Command-line interface example:
$  python3 -m tbb $your_script.py
Runs your_script.py in context of tbb.Monkey
"""

is_active = False
""" Indicates whether oneTBB context is activated """

ipc_enabled = False
""" Indicates whether IPC mode is enabled """

libirml = "libirml.so.1"


def _test(arg=None):
    """Some tests"""
    import platform
    if platform.system() == "Linux":
        ctypes.CDLL(libirml)
        assert 256 == os.system("ldd "+_api.__file__+"| grep -E 'libimf|libsvml|libintlc'") # nosec
    from .test import test
    test(arg)
    print("done")


def tbb_process_pool_worker27(inqueue, outqueue, initializer=None, initargs=(),
                            maxtasks=None):
    from multiprocessing.pool import worker
    worker(inqueue, outqueue, initializer, initargs, maxtasks)
    if ipc_enabled:
        try:
            librml = ctypes.CDLL(libirml)
            librml.release_resources()
        except:
            print("Warning: Can not load ", libirml, file=sys.stderr)


class TBBProcessPool27(multiprocessing.pool.Pool):
    def _repopulate_pool(self):
        """Bring the number of pool processes up to the specified number,
        for use after reaping workers which have exited.
        """
        from multiprocessing.util import debug

        for i in range(self._processes - len(self._pool)):
            w = self.Process(target=tbb_process_pool_worker27,
                             args=(self._inqueue, self._outqueue,
                                   self._initializer,
                                   self._initargs, self._maxtasksperchild)
                            )
            self._pool.append(w)
            w.name = w.name.replace('Process', 'PoolWorker')
            w.daemon = True
            w.start()
            debug('added worker')

    def __del__(self):
        self.close()
        for p in self._pool:
            p.join()

    def __exit__(self, *args):
        self.close()
        for p in self._pool:
            p.join()


def tbb_process_pool_worker3(inqueue, outqueue, initializer=None, initargs=(),
                            maxtasks=None, wrap_exception=False):
    from multiprocessing.pool import worker
    worker(inqueue, outqueue, initializer, initargs, maxtasks, wrap_exception)
    if ipc_enabled:
        try:
            librml = ctypes.CDLL(libirml)
            librml.release_resources()
        except:
            print("Warning: Can not load ", libirml, file=sys.stderr)


class TBBProcessPool3(multiprocessing.pool.Pool):
    def _repopulate_pool(self):
        """Bring the number of pool processes up to the specified number,
        for use after reaping workers which have exited.
        """
        from multiprocessing.util import debug

        for i in range(self._processes - len(self._pool)):
            w = self.Process(target=tbb_process_pool_worker3,
                             args=(self._inqueue, self._outqueue,
                                   self._initializer,
                                   self._initargs, self._maxtasksperchild,
                                   self._wrap_exception)
                            )
            self._pool.append(w)
            w.name = w.name.replace('Process', 'PoolWorker')
            w.daemon = True
            w.start()
            debug('added worker')

    def __del__(self):
        self.close()
        for p in self._pool:
            p.join()

    def __exit__(self, *args):
        self.close()
        for p in self._pool:
            p.join()


class Monkey:
    """
    Context manager which replaces standard multiprocessing.pool
    implementations with tbb.pool using monkey-patching. It also enables oneTBB
    threading for Intel(R) oneAPI Math Kernel Library (oneMKL). For example:

        with tbb.Monkey():
            run_my_numpy_code()

    It allows multiple parallel tasks to be executed on the same thread pool
    and coordinate number of threads across multiple processes thus avoiding
    overheads from oversubscription.
    """
    _items   = {}
    _modules = {}

    def __init__(self, max_num_threads=None, benchmark=False):
        """
        Create context manager for running under TBB scheduler.
        :param max_num_threads: if specified, limits maximal number of threads
        :param benchmark: if specified, blocks in initialization until requested number of threads are ready
        """
        if max_num_threads:
            self.ctl = global_control(global_control.max_allowed_parallelism, int(max_num_threads))
        if benchmark:
            if not max_num_threads:
               max_num_threads = default_num_threads()
            from .api import _concurrency_barrier
            _concurrency_barrier(int(max_num_threads))

    def _patch(self, class_name, module_name, obj):
        m = self._modules[class_name] = __import__(module_name, globals(),
                                                   locals(), [class_name])
        if m == None:
            return
        oldattr = getattr(m, class_name, None)
        if oldattr == None:
            self._modules[class_name] = None
            return
        self._items[class_name] = oldattr
        setattr(m, class_name, obj)

    def __enter__(self):
        global is_active
        assert is_active == False, "tbb.Monkey does not support nesting yet"
        is_active = True
        self.env_mkl = os.getenv('MKL_THREADING_LAYER')
        os.environ['MKL_THREADING_LAYER'] = 'TBB'
        self.env_numba = os.getenv('NUMBA_THREADING_LAYER')
        os.environ['NUMBA_THREADING_LAYER'] = 'TBB'

        if ipc_enabled:
            if sys.version_info.major == 2 and sys.version_info.minor >= 7:
                self._patch("Pool", "multiprocessing.pool", TBBProcessPool27)
            elif sys.version_info.major == 3 and sys.version_info.minor >= 5:
                self._patch("Pool", "multiprocessing.pool", TBBProcessPool3)
        self._patch("ThreadPool", "multiprocessing.pool", Pool)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        global is_active
        assert is_active == True, "modified?"
        is_active = False
        if self.env_mkl is None:
            del os.environ['MKL_THREADING_LAYER']
        else:
            os.environ['MKL_THREADING_LAYER'] = self.env_mkl
        if self.env_numba is None:
            del os.environ['NUMBA_THREADING_LAYER']
        else:
            os.environ['NUMBA_THREADING_LAYER'] = self.env_numba
        for name in self._items.keys():
            setattr(self._modules[name], name, self._items[name])


def init_sem_name():
    try:
        librml = ctypes.CDLL(libirml)
        librml.set_active_sem_name()
        librml.set_stop_sem_name()
    except Exception as e:
        print("Warning: Can not initialize name of shared semaphores:", e,
              file=sys.stderr)


def tbb_atexit():
    if ipc_enabled:
        try:
            librml = ctypes.CDLL(libirml)
            librml.release_semaphores()
        except:
            print("Warning: Can not release shared semaphores",
                  file=sys.stderr)


def _main():
    # Run the module specified as the next command line argument
    # python3 -m TBB user_app.py
    global ipc_enabled

    import platform
    import argparse
    parser = argparse.ArgumentParser(prog="python3 -m tbb", description="""
                Run your Python script in context of tbb.Monkey, which
                replaces standard Python pools and threading layer of
                Intel(R) oneAPI Math Kernel Library (oneMKL) by implementation based on
                Intel(R) oneAPI Threading Building Blocks (oneTBB). It enables multiple parallel
                tasks to be executed on the same thread pool and coordinate
                number of threads across multiple processes thus avoiding
                overheads from oversubscription.
             """, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    if platform.system() == "Linux":
        parser.add_argument('--ipc', action='store_true',
                        help="Enable inter-process (IPC) coordination between oneTBB schedulers")
        parser.add_argument('-a', '--allocator', action='store_true',
                        help="Enable oneTBB scalable allocator as a replacement for standard memory allocator")
        parser.add_argument('--allocator-huge-pages', action='store_true',
                        help="Enable huge pages for oneTBB allocator (implies: -a)")
    parser.add_argument('-p', '--max-num-threads', default=default_num_threads(), type=int,
                        help="Initialize oneTBB with P max number of threads per process", metavar='P')
    parser.add_argument('-b', '--benchmark', action='store_true',
                        help="Block oneTBB initialization until all the threads are created before continue the script. "
                        "This is necessary for performance benchmarks that want to exclude lazy scheduler initialization effects from the measurements")
    parser.add_argument('-v', '--verbose', action='store_true',
                        help="Request verbose and version information")
    parser.add_argument('-m', action='store_true', dest='module',
                        help="Executes following as a module")
    parser.add_argument('name', help="Script or module name")
    parser.add_argument('args', nargs=argparse.REMAINDER,
                        help="Command line arguments")
    args = parser.parse_args()

    if args.verbose:
        os.environ["TBB_VERSION"] = "1"
    if platform.system() == "Linux":
        if args.allocator_huge_pages:
            args.allocator = True
        if args.allocator and not os.environ.get("_TBB_MALLOC_PRELOAD"):
            libtbbmalloc_lib = 'libtbbmalloc_proxy.so.2'
            ld_preload = 'LD_PRELOAD'
            os.environ["_TBB_MALLOC_PRELOAD"] = "1"
            preload_list = filter(None, os.environ.get(ld_preload, "").split(':'))
            if libtbbmalloc_lib in preload_list:
                print('Info:', ld_preload, "contains", libtbbmalloc_lib, "already\n")
            else:
                os.environ[ld_preload] = ':'.join([libtbbmalloc_lib] + list(preload_list))

            if args.allocator_huge_pages:
                assert platform.system() == "Linux"
                try:
                    with open('/proc/sys/vm/nr_hugepages', 'r') as f:
                        pages = int(f.read())
                    if pages == 0:
                        print("oneTBB: Pre-allocated huge pages are not currently reserved in the system. To reserve, run e.g.:\n"
                              "\tsudo sh -c 'echo 2000 > /proc/sys/vm/nr_hugepages'")
                    os.environ["TBB_MALLOC_USE_HUGE_PAGES"] = "1"
                except:
                    print("oneTBB: Failed to read number of pages from /proc/sys/vm/nr_hugepages\n"
                          "\tIs the Linux kernel configured with the huge pages feature?")
                    sys.exit(1)

            os.execl(sys.executable, sys.executable, '-m', 'tbb', *sys.argv[1:])
            assert False, "Re-execution failed"

    sys.argv = [args.name] + args.args
    ipc_enabled = platform.system() == "Linux" and args.ipc
    os.environ["IPC_ENABLE"] = "1" if ipc_enabled else "0"
    if ipc_enabled:
        atexit.register(tbb_atexit)
        init_sem_name()
    if not os.environ.get("KMP_BLOCKTIME"): # TODO move
        os.environ["KMP_BLOCKTIME"] = "0"
    if '_' + args.name in globals():
        return globals()['_' + args.name](*args.args)
    else:
        import runpy
        runf = runpy.run_module if args.module else runpy.run_path
        with Monkey(max_num_threads=args.max_num_threads, benchmark=args.benchmark):
            runf(args.name, run_name='__main__')