# -*- coding: utf-8 -*-



from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import

from .logging   import *
from .task      import *
from .utils     import *

import collections
import time
import Queue
import multiprocessing
import signal
import timeit



class Runner(object):

    def __init__(self, db_server_uri, generic_task, task_args, context_args,
            workers_amount=0, period=1, log_server=None, log_collection=None,
            log_console=False, recover=False):

        assert generic_task
        assert task_args
        assert context_args

        # BUG: At this momment, this is commented due to the runner is in the
        # main process and logging module has a VERY BIG AND UGLY GLOBAL VAR

        # if log_server or log_console:
        #     logger_args = {
        #         "name":                         "recipes.runner",
        #     }
        #     if log_server:
        #         logger_args["db_uri"] =         log_server
        #         logger_args["db_collection"] =  log_collection or "recipes"
        #     if not log_console:
        #         logger_args["stream"]=          None
        #     get_default_logger(**logger_args)

        workers_amount = workers_amount or multiprocessing.cpu_count()
        to_process_queue = multiprocessing.Queue()
        processed_queue = multiprocessing.Queue()
        stop_event = multiprocessing.Event()
        stop_event.clear()
        clock_event = multiprocessing.Event()
        clock_event.clear()
        result = multiprocessing.Value('i', -1)

        # XXX: BEGIN TIMEIT
        start_time_1 = multiprocessing.Value("f", 0)
        start_time_2 = multiprocessing.Value("f", 0)
        stop_time = multiprocessing.Value("f", 0)
        # XXX: END TIMEIT

        self._stop_event = stop_event
        self._clock_event = clock_event
        self._result = result

        self._dispatcher = TaskDispatcher(db_server_uri, to_process_queue, processed_queue,
            clock_event, stop_event, generic_task, task_args, context_args,
            result, log_server, log_collection, log_console, recover,
            # XXX: BEGIN TIMEIT
            start_time_1, start_time_2, stop_time
            # XXX: END TIMEIT
        )

        self._timer = Timer(clock_event, stop_event, period, log_server,
            log_collection, log_console)

        self._workers = []
        for i in xrange(workers_amount):
            self._workers.append(TaskWorker(i+1, to_process_queue,
                processed_queue, clock_event, stop_event, log_server,
                log_collection, log_console,
                # XXX: BEGIN TIMEIT
                start_time_1, start_time_2, stop_time
                # XXX: END TIMEIT
            ))


    def start(self):
        start_time = timeit.default_timer()

        info("Starting Dispatcher")
        self._dispatcher.start()

        for w in self._workers:
            info("Starting Worker {}".format(w.id))
            w.start()

        info("Starting Timer")
        self._timer.start()


        info("Running...")
        while multiprocessing.active_children():
            try:
                self._clock_event.wait()
            except KeyboardInterrupt:
                warning("Keyboard interruption detected")
                warning("Warning children process")
                self._stop_event.set()
            except SystemExit as e:
                error("Interruption detected: " +  e.mmessage)
                info("Warning children process")
                self._stop_event.set()

        elapsed = timeit.default_timer() - start_time
        m, s = divmod(elapsed, 60)
        h, m = divmod(m, 60)
        info("Running time: {:02}:{:02}:{:02} ({})".format(int(h), int(m), int(s), elapsed))

        if self._result.value == 0:
            info("Result: OK")
        else:
            info("Result: ERROR (" + str(self._result.value) + ")")
        info("Quitting")

        return self._result.value



def _ignore_sigint():
    signal.signal(signal.SIGINT, signal.SIG_IGN)



class Timer(multiprocessing.Process):

    def __init__(self, clock_event, stop_event, period, log_server=None, log_collection=None, log_console=False):
        assert clock_event
        assert stop_event
        assert period > 0

        super(Timer, self).__init__()
        self.name = "TimerProcess"
        #self.daemon = True

        self._clock_event = clock_event
        self._stop_event = stop_event
        self._period = period
        self._log_server = log_server
        self._log_collection = log_collection
        self._log_console = log_console

    def run(self):
        if self._log_server or self._log_console:
            logger_args = {
                "name":                         "recipes.timer",
            }
            if self._log_server:
                logger_args["db_uri"] =         self._log_server
                logger_args["db_collection"] =  self._log_collection or "recipes"
            if not self._log_console:
                logger_args["stream"]=          None
            get_default_logger(**logger_args)
        _ignore_sigint()
        self._run()

    def _run(self):
        while not self._stop_event.is_set():
            try:
                self._clock_event.clear()
                time.sleep(self._period)
                self._clock_event.set()
            except (OSError, IOError):
                continue
        self._clock_event.set()
        info("Stopping")



class TaskState(object):

    REGISTERED =        1
    QUEUED =            2
    RUNNING =           3
    INPUT_FAILURE =     4
    RUN_SUCCESS =       5
    RUN_FAILURE =       6
    OUTPUT_FAILURE =    7

    STATES_STRINGS = {
        REGISTERED:     "Registered",
        QUEUED:         "Queued",
        RUNNING:        "Running",
        INPUT_FAILURE:  "Input location failure",
        RUN_SUCCESS:    "Execution success",
        RUN_FAILURE:    "Execution failure",
        OUTPUT_FAILURE: "Output location failure",
    }

    def __init__(self, specific_task):
        assert specific_task
        self._specific_task = specific_task
        self._dependencies = []
        self._state = self.REGISTERED
        self._valid_end = False
        self._aborted = False
        self.start_time = 0
        self.stop_time = 0

    @property
    def state(self):
        return self._state

    def valid_end():
        doc = "The valid_end property."
        def fget(self):
            return self._valid_end
        def fset(self, value):
            assert self.processed
            assert isinstance(value, bool)
            self._valid_end = value
        return locals()
    valid_end = property(**valid_end())

    def aborted():
        doc = "The aborted property."
        def fget(self):
            return self._aborted
        def fset(self, value):
            assert isinstance(value, bool)
            self._aborted = value
        return locals()
    aborted = property(**aborted())

    @property
    def registered(self):
        return self._state is self.REGISTERED

    @property
    def queued(self):
        return self._state is self.QUEUED

    @property
    def running(self):
        return self._state is self.RUNNING

    @property
    def input_failure(self):
        return self._state is self.INPUT_FAILURE

    @property
    def run_success(self):
        return self._state is self.RUN_SUCCESS

    @property
    def run_failure(self):
        return self._state is self.RUN_FAILURE

    @property
    def output_failure(self):
        return self._state is self.OUTPUT_FAILURE

    @property
    def processed(self):
        return self._state in [self.INPUT_FAILURE, self.RUN_SUCCESS,
            self.RUN_FAILURE, self.OUTPUT_FAILURE]

    @property
    def dependencies(self):
        return self._dependencies

    def depends_on(self, specific_task):
        assert specific_task
        if isinstance(specific_task, collections.Iterable):
            self._dependencies += specific_task
        else:
            self._dependencies.append(specific_task)

    def notify_registered(self):
        self._state = self.REGISTERED

    def notify_queued(self):
        assert self._state == self.REGISTERED
        self._state = self.QUEUED

    def notify_running(self):
        assert self._state == self.QUEUED
        self._state = self.RUNNING

    def notify_input_failure(self):
        assert self._state == self.QUEUED
        self._state = self.INPUT_FAILURE

    def notify_run_success(self):
        assert self._state == self.RUNNING
        self._state = self.RUN_SUCCESS

    def notify_run_failure(self):
        assert self._state == self.RUNNING
        self._state = self.RUN_FAILURE

    def notify_output_failure(self):
        assert self._state == self.RUNNING
        self._state = self.OUTPUT_FAILURE

    def textual_state(self):
        return self.STATES_STRINGS[self._state].upper()

    def __str__(self):
        parts = []
        if self._aborted:
            parts.append("(ABORTED)")
        if self._valid_end and self.processed and not self.run_success:
            parts.append("(VALID)")
        parts.append(self.textual_state())
        
        # XXX: Printing execution time of each task
        if self.processed:
            elapsed = self.stop_time - self.start_time
            m, s = divmod(elapsed, 60)
            h, m = divmod(m, 60)
            parts.append("{:02}:{:02}:{:02} ({})".format(int(h), int(m), int(s), elapsed))
        
        parts.append("->")
        parts.append(str(self._specific_task))
        return " ".join(parts)



class TaskRegister(collections.Mapping):

    def __init__(self, context_args):
        self._specific_tasks = {}
        self._context_args = context_args

    def __getitem__(self, key):
        return self._specific_tasks[key]

    def __setitem__(self, key, value):
        if key not in self._specific_tasks:
            raise KeyError(key)
        self._specific_tasks[key] = value

    def __iter__(self):
        return self._specific_tasks.iterkeys()

    def __len__(self):
        return len(self._specific_tasks)

    @property
    def has_remaining_tasks(self):
        return not all(s.aborted or s.processed
            for s in self._specific_tasks.itervalues())

    @property
    def valid_end(self):
        if self._specific_tasks:
            return all(s.valid_end for s in self._specific_tasks.itervalues())
        else:
            return False

    @property
    def dependencies_free_tasks(self):
        return (t for t in self._specific_tasks if self._is_free_task(t))

    def _is_free_task(self, specific_task):
        assert specific_task
        state = self._specific_tasks[specific_task]
        return (state.registered and not state.aborted and
            (not state.dependencies or
                all(self._specific_tasks[d].processed for d in state.dependencies)))

    def register(self, task, task_args=None):
        assert task
        specific_task = self._contextualize_task(task, task_args)
        try:
            task_state = self._specific_tasks[specific_task]
            new_register = False
        except KeyError:
            task_state = TaskState(specific_task)
            specific_deps = [self.register(d)[0] for d in ensure_iterable(specific_task.dependencies()) if d]
            if specific_deps:
                task_state.depends_on(specific_deps)
            self._specific_tasks[specific_task] = task_state
            new_register = True
        return specific_task, task_state, new_register

    def _contextualize_task(self, task, task_args=None):
        assert task
        if isinstance(task, type):
            task_args = task_args or {}
            # FIXME: Se puede quitar esta linea no?
            context = task.context_type(**self._context_args)
            task = task(**task_args)
        elif not isinstance(task, GenericTask):
            raise AssertionError("Unknown type")
        context = task.context_type(**self._context_args)
        task.set_context(context)
        task.init_context()
        return task

    def abort_tasks(self):
        for state in self._specific_tasks.itervalues():
            if not state.processed:
                state.aborted = True



from .recover   import *


class TaskDispatcher(multiprocessing.Process):

    def __init__(self, db_server_uri, to_process_queue, processed_queue,
            clock_event, stop_event, generic_task, task_args, context_args,
            result, log_server=None, log_collection=None, log_console=False,
            recover=False,
            # XXX: BEGIN TIMEIT
            start_time_1=None, start_time_2=None, stop_time=None
            # XXX: END TIMEIT
        ):
        assert db_server_uri
        assert to_process_queue
        assert processed_queue
        assert clock_event
        assert stop_event
        assert generic_task
        assert task_args
        assert context_args
        assert result

        super(TaskDispatcher, self).__init__()
        self.name = "DispatcherProcess"
        #self.daemon = True

        # XXX: BEGIN TIMEIT
        assert start_time_1
        assert start_time_2
        assert stop_time
        self._start_time_1 = start_time_1
        self._start_time_2 = start_time_2
        self._stop_time = stop_time
        # XXX: END TIMEIT

        self._db_server_uri = db_server_uri
        self._to_process_queue = to_process_queue
        self._processed_queue = processed_queue
        self._clock_event = clock_event
        self._stop_event = stop_event
        self._generic_task = generic_task
        self._task_args = task_args
        self._context_args = context_args
        self._log_server = log_server
        self._log_collection = log_collection
        self._log_console = log_console
        self._result = result
        self._recover = recover

    def run(self):
        if self._log_server or self._log_console:
            logger_args = {
                "name":                         "recipes.dispatcher",
            }
            if self._log_server:
                logger_args["db_uri"] =         self._log_server
                logger_args["db_collection"] =  self._log_collection or "recipes"
            if not self._log_console:
                logger_args["stream"]=          None
            get_default_logger(**logger_args)

        _ignore_sigint()
        self._run()

    def _run(self):
        try:
            info("Connecting to {}".format(self._db_server_uri))
            self._planned_tasks = RecoverableTaskRegister(
                self._db_server_uri, self._context_args, self._recover)
            info("Connection established")
            self._planned_tasks.register(self._generic_task, self._task_args)

            self._log_tasks(info, self._planned_tasks, "Tasks")

            while (not self._stop_event.is_set()
                    and self._planned_tasks.has_remaining_tasks):
                self._clock_event.wait()

                # self._log_tasks(debug, self._planned_tasks, "Tasks")

                self._handle_dependencies_free_tasks()
                self._handle_processed_tasks()

        except ValueError as e:
            import traceback
            error(e.message)
            traceback.print_exc()

        except Exception as e:
            import traceback
            fatal(e.__class__.__name__ + ": " + e.message)
            traceback.print_exc()

        # XXX: BEGIN TIMEIT
        else:
            with self._start_time_1.get_lock():
                assert self._start_time_1.value > 0
            with self._start_time_2.get_lock():
                assert self._start_time_2.value > 0
            with self._stop_time.get_lock():
                assert self._stop_time.value > 0

            with self._stop_time.get_lock(), self._start_time_1.get_lock():
                elapsed_1 = self._stop_time.value - self._start_time_1.value
            m, s = divmod(elapsed_1, 60)
            h, m = divmod(m, 60)
            txt = "Measured time 1: {:02}:{:02}:{:02} ({})".format(int(h), int(m), int(s), elapsed_1)
            info(txt)
            print("###")
            print("# " + txt)
            print("###")

            with self._stop_time.get_lock(), self._start_time_2.get_lock():
                elapsed_2 = self._stop_time.value - self._start_time_2.value
            m, s = divmod(elapsed_2, 60)
            h, m = divmod(m, 60)
            txt = "Measured time 2: {:02}:{:02}:{:02} ({})".format(int(h), int(m), int(s), elapsed_2)
            info(txt)
            print("###")
            print("# " + txt)
            print("###")
        # XXX: END TIMEIT

        finally:
            info("Stopping")
            self._stop_event.set()
            if self._planned_tasks.valid_end:
                info("All done")
                self._result.value = 0
            elif len(self._planned_tasks):
                warning("Not all tasks done")
                self._log_tasks(warning, self._planned_tasks, "Tasks")
                self._result.value = 1
            else:
                warning("No tasks")
                self._result.value = 2


    def _log_tasks(self, log_func, specific_tasks, title):
        assert log_func
        assert specific_tasks
        assert title
        log_func(title + ":")
        for specific_task in specific_tasks:
            state = self._planned_tasks[specific_task]
            log_func("")
            log_func("    " + str(state))
            for specific_dep in state.dependencies:
                log_func("        " + str(self._planned_tasks[specific_dep]))
            log_func("")

    def _handle_dependencies_free_tasks(self):
        for specific_task in self._planned_tasks.dependencies_free_tasks:
            try:
                self._planned_tasks[specific_task].notify_queued()
                self._to_process_queue.put_nowait(
                    (self._planned_tasks[specific_task], specific_task))
                self._log_tasks(info, [specific_task], "To process")
            except Queue.Full:
                warning("'To process' queue is full")
                break

    def _handle_processed_tasks(self):
        while True:
            try:
                state, specific_task = self._processed_queue.get_nowait()
                assert state.processed

                # Update state
                self._planned_tasks[specific_task] = state

                self._log_tasks(info, [specific_task], "Processed")
                self._pprint_state(state)

                if state.input_failure:
                    self._handle_failure(specific_task,
                        specific_task.on_incomplete_requirements)
                elif state.run_failure or state.output_failure:
                    self._handle_failure(specific_task,
                        specific_task.on_run_failure)
                elif state.run_success:
                    state.valid_end = True
                else:
                    raise AssertionError("Unknown state")
            except Queue.Empty:
                break

    def _pprint_state(self, state):
        task_done_str = str(state)
        if task_done_str[-2:] == "()":
            print(task_done_str)
        else:
            mark = " -> "
            offset = task_done_str.index(mark) + len(mark)
            state, task = task_done_str.split(" -> ")
            offset = "\n" + (offset * " ")
            task = task.replace("(", "(" + offset + "  ", 1)
            task = task.replace(", ", offset + "  ")
            task = task[:-1] + offset + ")"
            task = task.replace("=", " = ")
            new_task_done_str = " -> ".join([state, task])
            print(new_task_done_str)

    def _handle_failure(self, specific_task, specific_recovery_task_callback):
        assert specific_task

        if specific_recovery_task_callback:
            specific_recovery_task = specific_recovery_task_callback()
            if specific_recovery_task:
                self._abort_tasks()
                self._planned_tasks.register(specific_recovery_task)
            else:
                # Mark task as satisfactory because the function returns None in a explicit way
                self._planned_tasks[specific_task].valid_end = True
        else:
            warning("Not defined failure handler")
            self._abort_tasks()

    def _abort_tasks(self):
        self._planned_tasks.abort_tasks()
        while True:
            try:
                state, specific_task = self._to_process_queue.get_nowait()
                state.notify_registered()
            except Queue.Empty:
                break



class TaskWorker(multiprocessing.Process):

    def __init__(self, id, to_process_queue, processed_queue, clock_event, stop_event, log_server=None, log_collection=None, log_console=False,
        # XXX: BEGIN TIMEIT
        start_time_1=None, start_time_2=None, stop_time=None
        # XXX: END TIMEIT
    ):
        assert id > 0
        assert to_process_queue
        assert processed_queue
        assert clock_event
        assert stop_event

        super(TaskWorker, self).__init__()
        self.name = "Worker-" + str(id) + "Process"
        #self.daemon = True

        # XXX: BEGIN TIMEIT
        assert start_time_1
        assert start_time_2
        assert stop_time
        self._start_time_1 = start_time_1
        self._start_time_2 = start_time_2
        self._stop_time = stop_time
        # XXX: END TIMEIT

        self._id = id
        self._to_process_queue = to_process_queue
        self._processed_queue = processed_queue
        self._clock_event = clock_event
        self._stop_event = stop_event
        self._log_server = log_server
        self._log_collection = log_collection
        self._log_console = log_console

    @property
    def id(self):
        return self._id

    def run(self):
        if self._log_server or self._log_console:
            self._logger_args = {
                "name":                         "recipes.worker" + str(self.id),
            }
            if self._log_server:
                 self._logger_args["db_uri"] =         self._log_server
                 self._logger_args["db_collection"] =  self._log_collection or "recipes"
            if not self._log_console:
                 self._logger_args["stream"]=          None
            get_default_logger(**self._logger_args)

        _ignore_sigint()
        self._run()

    def _run(self):
        try:
            while not self._stop_event.is_set():
                self._clock_event.wait()
                try:
                    state, specific_task = self._to_process_queue.get_nowait()
                    assert state.queued

                    if not specific_task.context.args.dry_run:
                        try:
                            self._check_files(specific_task.input())
                        except TargetNotFoundError as e:
                            state.notify_input_failure()
                            if specific_task.on_incomplete_requirements:
                                warning(e.message)
                            else:
                                error(e.message)
                            assert state.processed
                            self._queue_processed_task(state, specific_task)
                            continue

                    self._dump_task(info, state, specific_task, "Executing")

                    try:
                        state.notify_running()
                        state.start_time = timeit.default_timer()
                        # XXX: BEGIN TIMEIT
                        if self._stop_time.value == 0 and specific_task.__class__.__name__ == "CreateFullDataset":
                            with self._stop_time.get_lock():
                                self._stop_time.value = state.start_time
                        elif self._start_time_1.value == 0 and specific_task.__class__.__name__ == "Instrument":
                            with self._start_time_1.get_lock():
                                self._start_time_1.value = state.start_time
                        elif self._start_time_2.value == 0 and specific_task.__class__.__name__ == "Dump":
                            with self._start_time_2.get_lock():
                                self._start_time_2.value = state.start_time
                        # XXX: END TIMEIT
                        specific_task.run()
                    except (ContextConfigurationError, ContextExecutionError) as e:
                        state.notify_run_failure()
                        info(e.message)
                        assert state.processed
                        self._queue_processed_task(state, specific_task)
                        continue
                    else:
                        info("Execution success")
                    finally:
                        state.stop_time = timeit.default_timer()


                    if not specific_task.context.args.dry_run:
                        try:
                            self._check_files(specific_task.output())
                        except TargetNotFoundError as e:
                            state.notify_output_failure()
                            if specific_task.on_run_failure:
                                warning(e.message)
                            else:
                                error(e.message)
                            assert state.processed
                            self._queue_processed_task(state, specific_task)
                            continue

                    state.notify_run_success()
                    assert state.processed
                    self._queue_processed_task(state, specific_task)
                except Queue.Empty:
                    continue
        except Exception as e:
            import traceback
            critical(traceback.format_exc())
            critical(e.__class__.__name__ + ": " + e.message)
        finally:
            warning("Warning other process")
            self._stop_event.set()
            info("Stopping")

    def _dump_task(self, log_func, state, specific_task, title):
        assert log_func
        assert state
        assert specific_task
        assert title

        log_func(title)
        log_func("")
        log_func("    {}".format(str(state)))
        log_func("")

    def _check_files(self, targets):
        if targets:
            if isinstance(targets, GenericTarget):
                targets.locate()
            elif isinstance(targets, collections.Mapping):
                self._check_files(targets.itervalues())
            elif isinstance(targets, collections.Iterable):
                [self._check_files(i) for i in targets]
            else:
                raise AssertionError("Unknown type of input")

    def _queue_processed_task(self, state, specific_task):
        assert state
        assert specific_task

        while True:
            try:
                self._processed_queue.put_nowait((state, specific_task))
                break
            except Queue.Full:
                warning("Done queue is full")
                debug("WORKER PUT HIT")
                self._clock_event.wait()
                continue
