from __future__ import absolute_import

import abc
import atexit
import copy
import logging
import os
import sys
import time
from collections import namedtuple

import gevent
import six

from dagster import ExecutionTargetHandle, PipelineDefinition, PipelineExecutionResult, check
from dagster.core.events import (
    DagsterEvent,
    DagsterEventType,
    PipelineProcessExitedData,
    PipelineProcessStartData,
    PipelineProcessStartedData,
)
from dagster.core.events.log import DagsterEventRecord
from dagster.core.execution.api import execute_run_iterator
from dagster.core.instance import DagsterInstance
from dagster.utils import get_multiprocessing_context
from dagster.utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info


class PipelineExecutionManager(six.with_metaclass(abc.ABCMeta)):
    @abc.abstractmethod
    def execute_pipeline(self, handle, pipeline, pipeline_run, instance, raise_on_error):
        '''Subclasses must implement this method.'''


def build_synthetic_pipeline_error_record(run_id, error_info, pipeline_name):
    check.str_param(run_id, 'run_id')
    check.str_param(pipeline_name, 'pipeline_name')
    check.inst_param(error_info, 'error_info', SerializableErrorInfo)

    return DagsterEventRecord(
        message=error_info.message + '\nStack Trace:\n' + '\n'.join(error_info.stack),
        # Currently it is the user_message that is displayed to the user client side
        # in dagit even though that was not the original intent. The original
        # intent was that the user_message was the message generated by user code
        # communicated directly to the client. We need to rationalize the treatment
        # of these different error messages
        user_message=(
            'An exception was thrown during execution that is likely a framework error, '
            'rather than an error in user code.'
        )
        + '\nOriginal error message: '
        + error_info.message
        + '\nStack Trace:\n'
        + '\n'.join(error_info.stack),
        level=logging.ERROR,
        run_id=run_id,
        timestamp=time.time(),
        error_info=error_info,
        pipeline_name=pipeline_name,
        dagster_event=DagsterEvent(DagsterEventType.PIPELINE_FAILURE.value, pipeline_name),
    )


def build_process_start_event(run_id, pipeline_name):
    check.str_param(pipeline_name, 'pipeline_name')
    check.str_param(run_id, 'run_id')
    message = 'About to start process for pipeline "{pipeline_name}" (run_id: {run_id}).'.format(
        pipeline_name=pipeline_name, run_id=run_id
    )

    return DagsterEventRecord(
        message=message,
        user_message=message,
        level=logging.INFO,
        run_id=run_id,
        timestamp=time.time(),
        error_info=None,
        pipeline_name=pipeline_name,
        dagster_event=DagsterEvent(
            message=message,
            event_type_value=DagsterEventType.PIPELINE_PROCESS_START.value,
            pipeline_name=pipeline_name,
            event_specific_data=PipelineProcessStartData(pipeline_name, run_id),
        ),
    )


def build_process_started_event(run_id, pipeline_name, process_id):
    message = 'Started process for pipeline (pid: {process_id}).'.format(process_id=process_id)

    return DagsterEventRecord(
        message=message,
        user_message=message,
        level=logging.INFO,
        run_id=run_id,
        timestamp=time.time(),
        error_info=None,
        pipeline_name=pipeline_name,
        dagster_event=DagsterEvent(
            message=message,
            event_type_value=DagsterEventType.PIPELINE_PROCESS_STARTED.value,
            pipeline_name=pipeline_name,
            step_key=None,
            solid_handle=None,
            step_kind_value=None,
            logging_tags=None,
            event_specific_data=PipelineProcessStartedData(
                pipeline_name=pipeline_name, run_id=run_id, process_id=process_id
            ),
        ),
    )


def build_process_exited_event(run_id, pipeline_name, process_id):
    message = 'Process for pipeline exited (pid: {process_id}).'.format(process_id=process_id)

    return DagsterEventRecord(
        message=message,
        user_message=message,
        level=logging.INFO,
        run_id=run_id,
        timestamp=time.time(),
        error_info=None,
        pipeline_name=pipeline_name,
        dagster_event=DagsterEvent(
            message=message,
            event_type_value=DagsterEventType.PIPELINE_PROCESS_EXITED.value,
            pipeline_name=pipeline_name,
            step_key=None,
            solid_handle=None,
            step_kind_value=None,
            logging_tags=None,
            event_specific_data=PipelineProcessExitedData(
                pipeline_name=pipeline_name, run_id=run_id, process_id=process_id
            ),
        ),
    )


class SynchronousExecutionManager(PipelineExecutionManager):
    def execute_pipeline(self, _, pipeline, pipeline_run, instance, raise_on_error):
        check.inst_param(pipeline, 'pipeline', PipelineDefinition)

        try:
            event_list = []
            for event in execute_run_iterator(pipeline, pipeline_run, instance):
                event_list.append(event)
            return PipelineExecutionResult(pipeline, pipeline_run.run_id, event_list, lambda: None)
        except Exception:  # pylint: disable=broad-except
            if raise_on_error:
                six.reraise(*sys.exc_info())

            instance.handle_new_event(
                build_synthetic_pipeline_error_record(
                    pipeline_run.run_id,
                    serializable_error_info_from_exc_info(sys.exc_info()),
                    pipeline.name,
                )
            )


class MultiprocessingExecutionManager(PipelineExecutionManager):
    def __init__(self):
        self._multiprocessing_context = get_multiprocessing_context()
        self._processes_lock = self._multiprocessing_context.Lock()
        self._processes = []
        # This is actually a reverse semaphore. We keep track of number of
        # processes we have by releasing semaphore every time we start
        # processing, we acquire after processing is finished
        self._processing_semaphore = gevent.lock.Semaphore(0)

        gevent.spawn(self._start_polling)
        atexit.register(self._cleanup)

    def _start_polling(self):
        while True:
            self._poll()
            gevent.sleep(0.1)

    def _cleanup(self):
        # Wait for child processes to finish and communicate on exit
        self.join()

    def _poll(self):
        with self._processes_lock:
            processes = copy.copy(self._processes)
            self._processes = []
            for _ in processes:
                self._processing_semaphore.release()

        for process in processes:
            done = self._consume_process_queue(process)
            if not done and not process.process.is_alive():
                done = self._consume_process_queue(process)
                if not done:
                    try:
                        done = True
                        raise Exception(
                            'Pipeline execution process for {run_id} unexpectedly exited'.format(
                                run_id=process.run_id
                            )
                        )
                    except Exception:  # pylint: disable=broad-except
                        process.instance.handle_new_event(
                            build_synthetic_pipeline_error_record(
                                process.run_id,
                                serializable_error_info_from_exc_info(sys.exc_info()),
                                process.pipeline_name,
                            )
                        )

            if not done:
                self._processes.append(process)

            self._processing_semaphore.acquire()

    def _consume_process_queue(self, process):
        while not process.message_queue.empty():
            message = process.message_queue.get(False)
            process.instance.handle_new_event(message)
            if isinstance(message, DagsterEventRecord) and (
                message.dagster_event.event_type_value
                == DagsterEventType.PIPELINE_PROCESS_EXITED.value
            ):
                return True
        return False

    def join(self):
        '''Waits until all there are no processes enqueued.'''
        while True:
            with self._processes_lock:
                if not self._processes and self._processing_semaphore.locked():
                    return True
            gevent.sleep(0.1)

    def execute_pipeline(self, handle, pipeline, pipeline_run, instance, raise_on_error):
        check.inst_param(handle, 'handle', ExecutionTargetHandle)
        check.invariant(
            raise_on_error is False, 'Multiprocessing execute_pipeline does not rethrow user error'
        )

        message_queue = self._multiprocessing_context.Queue()
        p = self._multiprocessing_context.Process(
            target=execute_pipeline_through_queue,
            kwargs={
                'handle': handle,
                'pipeline_run': pipeline_run,
                'instance_ref': instance.get_ref(),
                'message_queue': message_queue,
            },
        )

        instance.handle_new_event(build_process_start_event(pipeline_run.run_id, pipeline.name))

        p.start()
        with self._processes_lock:
            process = RunProcessWrapper(
                p, message_queue, instance, pipeline_run.run_id, pipeline.name
            )
            self._processes.append(process)


class RunProcessWrapper(
    namedtuple('RunProcessWrapper', 'process message_queue instance run_id pipeline_name')
):
    def __new__(cls, process, message_queue, instance, run_id, pipeline_name):
        return super(RunProcessWrapper, cls).__new__(
            cls, process, message_queue, instance, run_id, pipeline_name
        )


def execute_pipeline_through_queue(handle, pipeline_run, message_queue, instance_ref):
    """
    Execute pipeline using message queue as a transport
    """
    run_id = pipeline_run.run_id
    pipeline_name = pipeline_run.pipeline_name

    message_queue.put(build_process_started_event(run_id, pipeline_name, os.getpid()))

    try:
        handle.build_repository_definition()
        pipeline_def = handle.with_pipeline_name(pipeline_name).build_pipeline_definition()
    except Exception:  # pylint: disable=broad-except
        repo_error = sys.exc_info()
        message_queue.put(
            build_synthetic_pipeline_error_record(
                run_id, serializable_error_info_from_exc_info(repo_error), pipeline_name
            )
        )
        return

    try:
        event_list = []
        for event in execute_run_iterator(
            pipeline_def.build_sub_pipeline(pipeline_run.selector.solid_subset),
            pipeline_run,
            DagsterInstance.from_ref(instance_ref),
        ):
            event_list.append(event)
        return PipelineExecutionResult(pipeline_def, run_id, event_list, lambda: None)
    except Exception:  # pylint: disable=broad-except
        error_info = serializable_error_info_from_exc_info(sys.exc_info())
        message_queue.put(build_synthetic_pipeline_error_record(run_id, error_info, pipeline_name))
    finally:
        message_queue.put(build_process_exited_event(run_id, pipeline_name, os.getpid()))
        message_queue.close()
