
from __future__ import absolute_import

import collections
import functools
import logging
import threading

from vdsm.common import concurrent
from vdsm.common import exception
from vdsm.common import time


class NotRunning(Exception):
    """Executor not yet started or shutting down."""


class AlreadyStarted(Exception):
    """Executor started multiple times."""


class Executor(object):
    _log = logging.getLogger('Executor')

    def __init__(self, name, workers_count, max_tasks, scheduler,
                 max_workers=None, log=None):
        self._name = name
        self._workers_count = workers_count
        self._max_workers = max_workers
        self._worker_id = 0
        self._tasks = TaskQueue(name, max_tasks)
        self._scheduler = scheduler
        if log is not None:
            self._log = log
        self._workers = set()
        self._lock = threading.Lock()
        self._running = False

    def __repr__(self):
        return "<Executor %s workers=%d max_workers=%s %s at 0x%x>" % (
            self._name,
            self._workers_count,
            repr(self._tasks),
            id(self)
        )

    @property
    def name(self):
        return self._name

    def start(self):
        self._log.debug('Starting executor')
        with self._lock:
            if self._running:
                raise AlreadyStarted()
            self._running = True
            for _ in range(self._workers_count):
                self._add_worker()

    def stop(self, wait=True):
        self._log.debug('Stopping executor')
        with self._lock:
            self._running = False
            self._tasks.clear()
            for _ in range(self._workers_count):
                self._tasks.put(_STOP)
            workers = tuple(self._workers) if wait else ()
        for worker in workers:
            worker.join()

    def dispatch(self, callable, timeout=None, discard=True):
        if not self._running:
            raise NotRunning()
        self._tasks.put(Task(callable, timeout, discard))


    @property
    def _active_workers(self):
        return len([w for w in tuple(self._workers) if not w.discarded])

    @property
    def _total_workers(self):
        return len(self._workers)

    def _may_add_workers(self):
        return (self._active_workers < self._workers_count and
                (self._max_workers is None or
                 self._total_workers < self._max_workers))

    def _worker_discarded(self, worker):
        worker_added = False

        with self._lock:
            if not self._running:
                return
            if self._may_add_workers():
                self._add_worker()
                worker_added = True

        if not worker_added:
            self._log.warning("Too many workers (limit=%s), not adding more",
                              self._max_workers)
        self._log.warning("executor state: count=%d workers=%s",
                          self._total_workers, self._workers)

    def _worker_stopped(self, worker):
        worker_added = False

        with self._lock:
            self._workers.remove(worker)
            if not self._running:
                return
            if self._may_add_workers():
                self._add_worker()
                worker_added = True

        if worker_added:
            self._log.info("New worker added (%s active, %s total workers)",
                           self._active_workers, self._total_workers)

    def _next_task(self):
        task = self._tasks.get()
        if task is _STOP:
            raise NotRunning()
        return task


    def _add_worker(self):
        name = "%s/%d" % (self.name, self._worker_id)
        self._worker_id += 1
        worker = _Worker(self, self._scheduler, name, self._log)
        worker.start()
        self._workers.add(worker)


_STOP = object()


class _WorkerDiscarded(Exception):
    """ Raised if worker was discarded during execution of a task """


class _Worker(object):

    _log = logging.getLogger('Executor')

    def __init__(self, executor, scheduler, name, log=None):
        self._executor = executor
        self._scheduler = scheduler
        self._discarded = False
        self._task_counter = 0
        self._lock = threading.Lock()
        if log is not None:
            self._log = log
        self._thread = concurrent.thread(self._run, name=name, log=self._log)
        self._task = None
        self._scheduled_check = None

    @property
    def name(self):
        return self._thread.name

    def start(self):
        self._log.debug('Starting worker %s' % self.name)
        self._thread.start()

    def join(self):
        self._log.debug('Waiting for worker %s', self.name)
        self._thread.join()

    @property
    def discarded(self):
        return self._discarded

    def _run(self):
        self._log.debug('Worker started')
        try:
            while True:
                self._execute_task()
        except NotRunning:
            self._log.debug('Worker stopped')
        except _WorkerDiscarded:
            self._log.info('Worker was discarded')
        finally:
            self._executor._worker_stopped(self)

    def _execute_task(self):
        task = self._executor._next_task()
        with self._lock:
            self._scheduled_check = self._check_after(task.timeout)
        self._task = task
        try:
            task()
        except Exception:
            self._log.exception("Unhandled exception in %s", task)
        finally:
            self._task = None
            with self._lock:
                if self._scheduled_check is not None:
                    self._scheduled_check.cancel()
                    self._scheduled_check = None
                self._task_counter += 1
            if self._discarded:
                raise _WorkerDiscarded()

    def _check_after(self, timeout):
        if timeout is not None:
            check_task = functools.partial(
                self._check_task, self._task_counter)
            return self._scheduler.schedule(timeout, check_task)
        return None

    def _check_task(self, task_number):
        with self._lock:
            if task_number != self._task_counter:
                return
            if self._task.discard:
                if self._discarded:
                    raise AssertionError("Attempt to discard worker twice")
                self._discarded = True
            else:
                self._scheduled_check = self._check_after(self._task.timeout)
        if self._discarded:
            self._executor._worker_discarded(self)
            self._log.info("Worker discarded: %s", self)
        else:
            try:
                trace = concurrent.format_traceback(self._thread.ident)
            except KeyError:
                trace = "(traceback not available)"
            self._log.warning("Worker blocked: %s, traceback:\n%s", self,
                              trace)

    def __repr__(self):
            self.name,
            "running %s" % (self._task,) if self._task else "waiting",
            " discarded" if self._discarded else "",
            self._task_counter,
            id(self)
        )


class Task(object):

    def __init__(self, callable, timeout, discard=True):
        self._callable = callable
        self.timeout = timeout
        self.discard = discard
        self._start = None

    @property
    def duration(self):
        if self._start is None:
            return 0
        return time.monotonic_time() - self._start

    def __call__(self):
        self._start = time.monotonic_time()
        self._callable()

    def __repr__(self):
        return "<Task %s%s timeout=%s, duration=%.2f at 0x%x>" % (
            "discardable " if self.discard else "",
            self._callable,
            self.timeout,
            self.duration,
            id(self)
        )


class TaskQueue(object):

    def __init__(self, name, max_tasks):
        self._name = name
        self._max_tasks = max_tasks
        self._tasks = collections.deque()
        self._cond = threading.Condition(threading.Lock())

    def __repr__(self):
        return "<TaskQueue %s max_tasks=%i tasks(%i)=%s at 0x%x>" % (
            self._name,
            self._max_tasks,
            len(self._tasks),
            repr(self._tasks),
            id(self)
        )

    def put(self, task):
        with self._cond:
            if len(self._tasks) == self._max_tasks:
                raise exception.ResourceExhausted(
                    "Too many tasks",
                    resource=self._name,
                    current_tasks=self._max_tasks)
            self._tasks.append(task)
            self._cond.notify()

    def get(self):
        while True:
            try:
                return self._tasks.popleft()
            except IndexError:
                with self._cond:
                    if not self._tasks:
                        self._cond.wait()

    def clear(self):
        with self._cond:
            self._tasks.clear()
