"""RadicalPilotExecutor builds on the RADICAL-Pilot/Parsl
"""
import inspect
import logging
import os
import queue
import sys
import threading as mt
import time
from concurrent.futures import Future
from functools import partial
from pathlib import PosixPath
from typing import Dict, Optional

import requests
import typeguard

import parsl
from parsl.app.errors import BashExitFailure, RemoteExceptionWrapper
from parsl.app.python import timeout
from parsl.data_provider.files import File
from parsl.executors.base import ParslExecutor
from parsl.serialize import deserialize, pack_res_spec_apply_message
from parsl.serialize.errors import DeserializationError, SerializationError
from parsl.utils import RepresentationMixin

from .rpex_resources import CLIENT, MPI, ResourceConfig

try:
    import radical.pilot as rp
    import radical.utils as ru
except ImportError:
    _rp_enabled = False
else:
    _rp_enabled = True


RPEX = 'RPEX'
BASH = 'bash'
PYTHON = 'python'

CWD = os.getcwd()
PWD = os.path.abspath(os.path.dirname(__file__))

PARSL_RP_RESOURCE_MAP = {'cores': 'ranks',
                         'disk': 'lfs_per_rank',
                         'memory': 'mem_per_rank'}

logger = logging.getLogger(__name__)


class RadicalPilotExecutor(ParslExecutor, RepresentationMixin):
    """Executor is designed for executing heterogeneous tasks
       in terms of type/resource.

    The RadicalPilotExecutor system has the following main components:

      1. "start": Create and start the RADICAL-Pilot runtime components ``rp.Session``,
          ``rp.PilotManager`` and ``rp.TaskManager``.
      2. "translate": Unwrap, identify, and parse Parsl ``apps`` into ``rp.TaskDescription``.
      3. "submit": Submit Parsl apps to ``rp.TaskManager``.
      4. "shutdown": Shut down the RADICAL-Pilot runtime and all associated components.

    Here is a diagram

    .. code:: python

      ----------------------------------------------------------------------------
               Parsl Data Flow Kernel        |   Task Translator |  rp.TaskManager
      ---------------------------------------|-------------------|----------------
                                             |                   |
      -> Dep. check ------> Parsl_tasks{} <--+--> Parsl Task     |
       Data management          +dfk.submit  |        |          |
                                             |        v          |
                                             |    RP Task(s) ->  | submit(task)
      ----------------------------------------------------------------------------

    The RadicalPilotExecutor creates a ``rp.Session``, ``rp.TaskManager``,
    and ``rp.PilotManager``. The executor receives the Parsl apps from the
    DFK and translates these apps (in-memory) into ``rp.TaskDescription``
    object to be passed to the ``rp.TaskManager``. This executor has two
    submission mechanisms:

    1. Default_mode: where the executor submits the tasks directly to
       RADICAL-Pilot.

    2. Bulk_mode: where the executor accumulates N tasks (functions and
       executables) and submit them.

    Parameters
    ----------
    rpex_cfg : :class: `~parsl.executors.rpex_resources.ResourceConfig`
        a dataclass specifying resource configuration.
        Default is ResourceConfig instance.

    label : str
        Label for this executor instance.
        Default is "RPEX".

    bulk_mode : bool
        Enable bulk mode submission and execution. Default is False (stream).

    resource : Optional[str]
        The resource name of the targeted HPC machine or cluster.
        Default is local.localhost (user local machine).

    runtime : int
        The maximum runtime for the entire job in minutes.
        Default is 30.

    working_dir : str
        The working dir to be used by the executor.

    rpex_pilot_kwargs: Dict of kwargs that are passed directly to the rp.PilotDescription object.

    For more information: https://radicalpilot.readthedocs.io/en/stable/
    """

    @typeguard.typechecked
    def __init__(self,
                 resource: str,
                 label: str = RPEX,
                 bulk_mode: bool = False,
                 working_dir: Optional[str] = None,
                 rpex_cfg: Optional[ResourceConfig] = None, **rpex_pilot_kwargs):

        super().__init__()
        self.pmgr = None
        self.tmgr = None
        self.run_dir = '.'
        self.label = label
        self.session = None
        self.resource = resource
        self._uid = RPEX.lower()
        self.bulk_mode = bulk_mode
        self._terminate = mt.Event()
        self.working_dir = working_dir
        self.pilot_kwargs = rpex_pilot_kwargs
        self.future_tasks: Dict[str, Future] = {}

        if rpex_cfg:
            self.rpex_cfg = rpex_cfg.get_config()
        elif not rpex_cfg and 'local' in resource:
            self.rpex_cfg = ResourceConfig().get_config()
        else:
            raise ValueError('Resource config must be '
                             'specified for a non-local resources')

    def task_state_cb(self, task, state):
        """
        Update the state of Parsl Future apps
        Based on RP task state callbacks.
        """
        # check the Master/Worker state
        if task.mode in [rp.RAPTOR_MASTER, rp.RAPTOR_WORKER]:
            if state == rp.FAILED:
                exception = RuntimeError(f'{task.uid} failed with internal error: {task.stderr}')
                self._fail_all_tasks(exception)

        # check all other tasks state
        else:
            parsl_task = self.future_tasks[task.uid]

            if state == rp.DONE:
                if task.description['mode'] in [rp.TASK_EXEC,
                                                rp.TASK_PROC,
                                                rp.TASK_EXECUTABLE]:
                    parsl_task.set_result(int(task.exit_code))
                else:
                    # we do not support MPI function output
                    # serialization. TODO: To be fixed soon.
                    if not task.description.get('use_mpi'):
                        result = deserialize(eval(task.return_value))
                        parsl_task.set_result(result)
                    else:
                        parsl_task.set_result(task.return_value)

            elif state == rp.CANCELED:
                parsl_task.cancel()

            elif state == rp.FAILED:
                if task.description['mode'] in [rp.TASK_EXEC,
                                                rp.TASK_EXECUTABLE]:
                    parsl_task.set_exception(BashExitFailure(task.name,
                                                             task.exit_code))
                else:
                    if task.exception:
                        # unpack a serialized exception
                        if not task.description.get('use_mpi') or task.description['mode'] == rp.TASK_PROC:
                            self._unpack_and_set_parsl_exception(parsl_task, task.exception)
                        # we do not serialize mpi function exception
                        else:
                            parsl_task.set_exception(eval(task.exception))
                    else:
                        parsl_task.set_exception('Task failed for an unknown reason')

    def _fail_all_tasks(self, exception):
        """
        Fail all outstanding tasks with the given exception.

        This method iterates through all outstanding tasks in the
        `_future_tasks` dictionary, which have not yet completed,
        and sets the provided exception as their result, indicating
        a failure.

        Parameters:
        - exception: The exception to be set as the result for all
                     outstanding tasks.
        """
        for fut_task in self.future_tasks.values():
            if not fut_task.done():
                fut_task.set_exception(exception)

    def start(self):
        """Create the Pilot component and pass it.
        """
        logger.info("starting RadicalPilotExecutor")
        logger.info('Parsl: {0}'.format(parsl.__version__))
        logger.info('RADICAL pilot: {0}'.format(rp.version))
        self.session = rp.Session(cfg={'base': self.run_dir},
                                  uid=ru.generate_id('rpex.session',
                                                     mode=ru.ID_PRIVATE))
        logger.info("RPEX session is created: {0}".format(self.session.path))

        pd_init = {**self.pilot_kwargs,
                   'exit_on_error': True,
                   'resource': self.resource}

        if not self.resource or 'local' in self.resource:
            os.environ["RADICAL_LOG_LVL"] = "DEBUG"
            logger.info("RPEX will be running in local mode")

        pd = rp.PilotDescription(pd_init)
        pd.verify()

        # start RP's main components TMGR, PMGR and Pilot
        self.tmgr = rp.TaskManager(session=self.session)
        self.pmgr = rp.PilotManager(session=self.session)
        self.pilot = self.pmgr.submit_pilots(pd)

        if not self.pilot.description.get('cores') or not self.pilot.description.get('nodes'):
            logger.warning('no "cores/nodes" per pilot were set, using default resources')

        self.tmgr.add_pilots(self.pilot)
        self.tmgr.register_callback(self.task_state_cb)

        worker_path = '{0}/rpex_worker.py'.format(PWD)

        self.masters = []

        logger.info(f'Starting {self.rpex_cfg.n_masters} masters and {self.rpex_cfg.n_workers} workers for each master')

        # create N masters
        for _ in range(self.rpex_cfg.n_masters):
            md = rp.TaskDescription(self.rpex_cfg.master_descr)
            md.uid = ru.generate_id('rpex.master.%(item_counter)06d', ru.ID_CUSTOM,
                                    ns=self.session.uid)

            # submit the master to the TMGR
            master = self.tmgr.submit_raptors(md)[0]
            self.masters.append(master)

            workers = []
            # create N workers for each master and submit them to the TMGR
            for _ in range(self.rpex_cfg.n_workers):
                wd = rp.TaskDescription(self.rpex_cfg.worker_descr)
                wd.uid = ru.generate_id('rpex.worker.%(item_counter)06d', ru.ID_CUSTOM,
                                        ns=self.session.uid)
                wd.raptor_id = master.uid
                wd.input_staging = self._stage_files([File(worker_path)], mode='in')
                workers.append(wd)

            self.tmgr.submit_workers(workers)

        self.select_master = self._cyclic_master_selector()

        # prepare or use the current env for the agent/pilot side environment
        if self.rpex_cfg.pilot_env_mode != CLIENT:
            logger.info("creating {0} environment for the executor".format(self.rpex_cfg.pilot_env.name))
            self.pilot.prepare_env(env_name=self.rpex_cfg.pilot_env.name,
                                   env_spec=self.rpex_cfg.pilot_env.as_dict())
        else:
            client_env = sys.prefix
            logger.info("reusing ({0}) environment for the executor".format(client_env))

        # create a bulking thread to run the actual task submission
        # to RP in bulks
        if self.bulk_mode:
            self._max_bulk_size = 1024
            self._max_bulk_time = 3        # seconds
            self._min_bulk_time = 0.1      # seconds

            self._bulk_queue = queue.Queue()
            self._bulk_thread = mt.Thread(target=self._bulk_collector)

            self._bulk_thread.daemon = True
            self._bulk_thread.start()

            logger.info('bulk mode is on, submitting tasks in bulks')

        return True

    def _cyclic_master_selector(self):
        """
        Balance tasks submission across N masters and N workers
        """
        current_master = 0
        masters_uids = [m.uid for m in self.masters]

        while True:
            yield masters_uids[current_master]
            current_master = (current_master + 1) % len(self.masters)

    def unwrap(self, func, args):
        """
        Unwrap a Parsl app and its args for further processing.

        Parameters
        ----------
        func : callable
            The function to be unwrapped.

        args : tuple
            The arguments associated with the function.

        Returns
        -------
        tuple
            A tuple containing the unwrapped function, adjusted arguments,
            and task type information.
        """

        task_type = ''

        while hasattr(func, '__wrapped__'):
            func = func.__wrapped__

        try:
            if isinstance(func, partial):
                try:
                    task_type = inspect.getsource(func.args[0]).split('\n')[0]
                    if BASH in task_type:
                        task_type = BASH
                        func = func.args[0]
                    else:
                        task_type = PYTHON

                except Exception:
                    logger.exception('unwrap failed')

                    return func, args, task_type

            else:
                task_type = inspect.getsource(func).split('\n')[0]
                if PYTHON in task_type:
                    task_type = PYTHON
                else:
                    task_type = ''
        except Exception as e:
            raise Exception('failed to obtain task type: {0}'.format(e))

        return func, args, task_type

    def task_translate(self, tid, func, parsl_resource_specification, args, kwargs):
        """
        Convert Parsl function to RADICAL-Pilot rp.TaskDescription
        """

        task = rp.TaskDescription()
        task.name = func.__name__

        if parsl_resource_specification and isinstance(parsl_resource_specification, dict):
            logger.debug('mapping Parsl resource specifications >> RP resource specifications')
            for key, val in parsl_resource_specification.items():
                if key not in task.as_dict():
                    key = PARSL_RP_RESOURCE_MAP.get(key, None)
                    if not key:
                        logger.warning('ignoring "{0}" key from task resource specification as it is not supported by RP'.format(key))
                        continue
                setattr(task, key, val)

        func, args, task_type = self.unwrap(func, args)

        if BASH in task_type:
            if callable(func):
                # if the user specifies the executable mode then we expect the
                # a code in a file that need to be executed in an isolated env.
                if parsl_resource_specification.get('mode') == rp.TASK_EXECUTABLE:
                    # These lines of code are from parsl/app/bash.py
                    try:
                        # Execute the func to get the command
                        bash_app = func(*args, **kwargs)
                        if not isinstance(bash_app, str):
                            raise ValueError("Expected a str for bash_app cmd,"
                                             "got: {0}".format(type(bash_app)))
                    except AttributeError as e:
                        raise Exception("failed to obtain bash app cmd") from e

                    task.executable = bash_app
                    task.mode = rp.TASK_EXECUTABLE

                # This is the default mode where the bash_app will be executed as
                # as a single core process by RP. For cores > 1 the user must use
                # task.mode=rp.TASK_EXECUTABLE (above) or use MPI functions if their
                # code is Python.
                else:
                    task.mode = rp.TASK_PROC
                    task.raptor_id = next(self.select_master)
                    task.executable = self._pack_and_apply_message(func, args, kwargs)

        elif PYTHON in task_type or not task_type:
            task.mode = rp.TASK_FUNCTION
            task.raptor_id = next(self.select_master)
            if kwargs.get('walltime'):
                func = timeout(func, kwargs['walltime'])

            # Check how to serialize the function object
            if MPI in self.rpex_cfg.worker_type.lower():
                task.use_mpi = True
                task.function = rp.PythonTask(func, *args, **kwargs)
            else:
                task.use_mpi = False
                task.function = self._pack_and_apply_message(func, args, kwargs)

        task.input_staging = self._stage_files(kwargs.get("inputs", []),
                                               mode='in')
        task.output_staging = self._stage_files(kwargs.get("outputs", []),
                                                mode='out')

        task.input_staging.extend(self._stage_files(list(args), mode='in'))

        self._set_stdout_stderr(task, kwargs)

        try:
            task.verify()
        except ru.typeddict.TDKeyError as e:
            raise Exception(f'{e}. Please check: https://radicalpilot.readthedocs.io/en/stable/ documentation')

        return task

    def _pack_and_apply_message(self, func, args, kwargs):
        try:
            buffer = pack_res_spec_apply_message(func,
                                                 args,
                                                 kwargs,
                                                 resource_specification={},
                                                 buffer_threshold=1024 * 1024)
            task_func = rp.utils.serialize_bson(buffer)
        except TypeError:
            raise SerializationError(func.__name__)

        return task_func

    def _unpack_and_set_parsl_exception(self, parsl_task, exception):
        try:
            try:
                s = rp.utils.deserialize_bson(exception)
            except Exception:
                s = exception

            if isinstance(s, RemoteExceptionWrapper):
                try:
                    s.reraise()
                except Exception as e:
                    parsl_task.set_exception(e)
            elif isinstance(s, Exception):
                parsl_task.set_exception(s)
            elif isinstance(s, str):
                parsl_task.set_exception(eval(s))
            else:
                raise ValueError("Unknown exception-like type received: {}".format(type(s)))
        except Exception as e:
            parsl_task.set_exception(
                DeserializationError("Received exception, but handling also threw an exception: {}".format(e)))

    def _set_stdout_stderr(self, task, kwargs):
        """
        set the stdout and stderr of a task
        """
        for k in ['stdout', 'stderr']:
            k_val = kwargs.get(k, '')
            if k_val:
                # check the type of the stderr/out
                if isinstance(k_val, File):
                    k_val = k_val.filepath
                elif isinstance(k_val, PosixPath):
                    k_val = k_val.__str__()

                # set the stderr/out to the desired
                # name by the user
                setattr(task, k, k_val)
                task.sandbox = CWD

    def _stage_files(self, files, mode):
        """
        a function to stage list of input/output
        files between two locations.
        """
        to_stage = []
        files = [f for f in files if isinstance(f, File)]
        for file in files:
            if mode == 'in':
                # a workaround RP not supporting
                # staging https file
                if file.scheme == 'https':
                    r = requests.get(file.url)
                    p = CWD + '/' + file.filename
                    with open(p, 'wb') as ff:
                        ff.write(r.content)
                    file = File(p)

                f = {'source': file.url,
                     'action': rp.TRANSFER}
                to_stage.append(f)

            elif mode == 'out':
                # this indicates that the user
                # did not provided a specific
                # output file and RP will stage out
                # the task.stdout from pilot://task_folder
                # to the CWD or file.url
                if '/' not in file.url:
                    f = {'source': file.filename,
                         'target': file.url,
                         'action': rp.TRANSFER}
                    to_stage.append(f)
            else:
                raise ValueError('unknown staging mode')

        return to_stage

    def _bulk_collector(self):

        bulk = list()

        while not self._terminate.is_set():

            now = time.time()  # time of last submission

            # collect tasks for min bulk time
            # NOTE: total collect time could actually be max_time + min_time
            while time.time() - now < self._max_bulk_time:

                try:
                    task = self._bulk_queue.get(block=True,
                                                timeout=self._min_bulk_time)
                except queue.Empty:
                    task = None

                if task:
                    bulk.append(task)

                if len(bulk) >= self._max_bulk_size:
                    break

                if self._terminate.is_set():
                    break

            if bulk:
                logger.debug('submit bulk: %d', len(bulk))
                self.tmgr.submit_tasks(bulk)
                bulk = list()

    def submit(self, func, resource_specification, *args, **kwargs):
        """
        Submits tasks in stream mode or bulks (bulk mode)
        to RADICAL-Pilot rp.TaskManager.
        """
        rp_tid = ru.generate_id('task.%(item_counter)06d', ru.ID_CUSTOM,
                                ns=self.session.uid)
        parsl_tid = int(rp_tid.split('task.')[1])

        logger.debug("got Task {0} from Parsl-dfk".format(parsl_tid))
        task = self.task_translate(parsl_tid, func, resource_specification, args, kwargs)

        # assign task id for rp task
        task.uid = rp_tid

        # set the future with corresponding id
        self.future_tasks[rp_tid] = Future()

        if self.bulk_mode:
            # push task to rp submit thread
            self._bulk_queue.put(task)
        else:
            # submit the task to rp
            logger.debug("put {0} to rp-TMGR".format(rp_tid))
            self.tmgr.submit_tasks(task)

        return self.future_tasks[rp_tid]

    def shutdown(self, hub=True, targets='all', block=False):
        """Shutdown the executor, including all RADICAL-Pilot components."""
        logger.info("RadicalPilotExecutor is terminating...")

        self._terminate.set()

        # ensure we are in the bulk submssion mode
        if self.bulk_mode:
            self._bulk_thread.join()

        self.session.close(download=True)
        logger.info("RadicalPilotExecutor is terminated.")

        return True