"""Implement facilities for executing asynchronous tasks."""
import logging

from threading import Event, Lock, RLock
from types import TracebackType
from typing import Any, Callable, Optional, Tuple, Type

from . import thread_pool
from .postgresql.cancellable import CancellableSubprocess

logger = logging.getLogger(__name__)


class CriticalTask(object):
    """Represents a critical task in a background process that we either need to cancel or get the result of.

    Fields of this object may be accessed only when holding a lock on it. To perform the critical task the background
    thread must, while holding lock on this object, check ``is_cancelled`` flag, run the task and mark the task as
    complete using :func:`complete`.

    The main thread must hold async lock to prevent the task from completing, hold lock on critical task object,
    call :func:`cancel`. If the task has completed :func:`cancel` will return ``False`` and ``result`` field will
    contain the result of the task. When :func:`cancel` returns ``True`` it is guaranteed that the background task will
    notice the ``is_cancelled`` flag.

    :ivar is_cancelled: if the critical task has been cancelled.
    :ivar result: contains the result of the task, if it has already been completed.
    """

    def __init__(self) -> None:
        """Create a new instance of :class:`CriticalTask`.

        Instantiate the lock and the task control attributes.
        """
        self._lock = Lock()
        self.is_cancelled = False
        self.result = None

    def reset(self) -> None:
        """Must be called every time the background task is finished.

        .. note::
            Must be called from async thread. Caller must hold lock on async executor when calling.
        """
        self.is_cancelled = False
        self.result = None

    def cancel(self) -> bool:
        """Tries to cancel the task.

        .. note::
            Caller must hold lock on async executor and the task when calling.

        :returns: ``False`` if the task has already run, or ``True`` it has been cancelled.
        """
        if self.result is not None:
            return False
        self.is_cancelled = True
        return True

    def complete(self, result: Any) -> None:
        """Mark task as completed along with a *result*.

        .. note::
            Must be called from async thread. Caller must hold lock on task when calling.
        """
        self.result = result

    def __enter__(self) -> 'CriticalTask':
        """Acquire the object lock when entering the context manager."""
        self._lock.acquire()
        return self

    def __exit__(self, exc_type: Optional[Type[BaseException]],
                 exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]) -> None:
        """Release the object lock when exiting the context manager."""
        self._lock.release()


class AsyncExecutor(object):
    """Asynchronous executor of (long) tasks.

    :ivar critical_task: a :class:`CriticalTask` instance to handle execution of critical background tasks.
    """

    def __init__(self, cancellable: CancellableSubprocess, ha_wakeup: Callable[..., None]) -> None:
        """Create a new instance of :class:`AsyncExecutor`.

        Configure the given *cancellable* and *ha_wakeup*, initializes the control attributes, and instantiate the lock
        and event objects that are used to access attributes and manage communication between threads.

        :param cancellable: a subprocess that supports being cancelled.
        :param ha_wakeup: function to wake up the HA loop.
        """
        self._cancellable = cancellable
        self._ha_wakeup = ha_wakeup
        self._thread_lock = RLock()
        self._scheduled_action: Optional[str] = None
        self._scheduled_action_lock = RLock()
        self._is_cancelled = False
        self._finish_event = Event()
        self.critical_task = CriticalTask()

    @property
    def busy(self) -> bool:
        """``True`` if there is an action scheduled to occur, else ``False``."""
        return self.scheduled_action is not None

    def schedule(self, action: str) -> Optional[str]:
        """Schedule *action* to be executed.

        .. note::
            Must be called before executing a task.

        .. note::
            *action* can only be scheduled if there is no other action currently scheduled.

        :param action: action to be executed.

        :returns: ``None`` if *action* has been successfully scheduled, or the previously scheduled action, if any.
        """
        with self._scheduled_action_lock:
            if self._scheduled_action is not None:
                return self._scheduled_action
            self._scheduled_action = action
            self._is_cancelled = False
            self._finish_event.set()
        return None

    @property
    def scheduled_action(self) -> Optional[str]:
        """The currently scheduled action, if any, else ``None``."""
        with self._scheduled_action_lock:
            return self._scheduled_action

    def reset_scheduled_action(self) -> None:
        """Unschedule a previously scheduled action, if any.

        .. note::
            Must be called once the scheduled task finishes or is cancelled.
        """
        with self._scheduled_action_lock:
            self._scheduled_action = None

    def run(self, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> Optional[Any]:
        """Run *func* with *args*.

        .. note::
            Expected to be executed through a thread.

        :param func: function to be run. If it returns anything other than ``None``, HA loop will be woken up at the end
            of :func:`run` execution.
        :param args: arguments to be passed to *func*.

        :returns: ``None`` if *func* execution has been cancelled or faced any exception, otherwise the result of
            *func*.
        """
        wakeup = False
        try:
            with self:
                if self._is_cancelled:
                    return
                self._finish_event.clear()

            self._cancellable.reset_is_cancelled()
            # if the func returned something (not None) - wake up main HA loop
            wakeup = func(*args) if args else func()
            return wakeup
        except Exception:
            logger.exception('Exception during execution of long running task %s', self.scheduled_action)
        finally:
            with self:
                self.reset_scheduled_action()
                self._finish_event.set()
                with self.critical_task:
                    self.critical_task.reset()
            if wakeup is not None:
                self._ha_wakeup()

    def run_async(self, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> None:
        """Start an async thread that runs *func* with *args*.

        :param func: function to be run.
        :param args: arguments to be passed along with *func*.
        """
        thread_pool.get_executor().submit(self.run, func, args)

    def try_run_async(self, action: str, func: Callable[..., Any], args: Tuple[Any, ...] = ()) -> Optional[str]:
        """Try to run an async task, if none is currently being executed.

        :param action: name of the task to be executed.
        :param func: actual function that performs the task *action*.
        :param args: arguments to be passed to *func*.

        :returns: ``None`` if *func* was scheduled successfully, otherwise an error message informing of an already
            ongoing task.
        """
        prev = self.schedule(action)
        if prev is None:
            return self.run_async(func, args)
        return 'Failed to run {0}, {1} is already in progress'.format(action, prev)

    def cancel(self) -> None:
        """Request cancellation of a scheduled async task, if any.

        .. note::
            Wait until task is cancelled before returning control to caller.
        """
        with self:
            with self._scheduled_action_lock:
                if self._scheduled_action is None:
                    return
                logger.warning('Cancelling long running task %s', self._scheduled_action)
            self._is_cancelled = True

        self._cancellable.cancel()
        self._finish_event.wait()

        with self:
            self.reset_scheduled_action()

    def __enter__(self) -> 'AsyncExecutor':
        """Acquire the thread lock when entering the context manager."""
        self._thread_lock.acquire()
        return self

    def __exit__(self, exc_type: Optional[Type[BaseException]],
                 exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]) -> None:
        """Release the thread lock when exiting the context manager.

        .. note::
            The arguments are not used, but we need them to match the expected method signature.
        """
        self._thread_lock.release()
