import asyncio
import time
from concurrent.futures.thread import ThreadPoolExecutor
from contextlib import asynccontextmanager
from functools import partial
from typing import Dict, AsyncContextManager, Callable, Set, Sequence, Optional

from clarabot_nano_common.utility.path.tracker import PathDependencyTracker

from clarabot_nano.common.filesystem.operation.base import AbstractIOQueueOperation


class TokenLock:
    __slots__ = ("path", "token", "until")

    def __init__(self, path: str, token: str, until: float):
        self.path = path
        self.token = token
        self.until = until


class TokenLockError(Exception):
    def __init__(self, token: str):
        Exception.__init__(self, token)

    @property
    def token(self) -> str:
        return self.args[0]


# TODO: add persistence to ex-locks for more consistent behaviour
class FilesystemTxQueue(object):
    DEFAULT_TOKEN_LOCK_TTL = 30 * 60

    def __init__(self, pool: ThreadPoolExecutor):
        self.pool = pool
        # track pending operations by account so they can be cancelled later
        self._ops_by_account: Dict[int, Set[AbstractIOQueueOperation]] = dict()
        # map blocking path to future; single for blocking exclusion
        self._locks: Dict[str, asyncio.Future] = dict()
        self._tracker = PathDependencyTracker()
        # token_lock
        self._token_locks: Dict[str, TokenLock] = dict()
        self._token_tracker = PathDependencyTracker()

    async def close(self):
        futures = set()
        for f in self._locks.values():
            f.cancel()
            futures.add(f)
        await asyncio.gather(*list(futures), return_exceptions=True)

    @asynccontextmanager
    async def wait_ready(
        self,
        operation: Optional[AbstractIOQueueOperation] = None,
        paths: Optional[Sequence[str]] = None,
        break_ex_locks: bool = False,
    ) -> AsyncContextManager[Callable[[], asyncio.Future]]:
        if not paths:
            paths = operation.paths()
        elif operation:
            raise RuntimeError("Do not provide operation and paths both.")
        # Wait for blocking operation to complete (or until cancellation).
        while True:
            blocking_futures = set()
            for path in paths:
                for affecting_path in self._tracker.affects(path):
                    blocking_futures.add(self._locks[affecting_path])
            if not blocking_futures:
                # Break explicit locks by tokens or raise if they lock the target path(s).
                if break_ex_locks:
                    for path in paths:
                        self._break_ex_locks(path)
                else:
                    for path in paths:
                        lock = self._get_ex_lock_best(path)
                        if lock is not None:
                            raise TokenLockError(lock.token)
                break
            await asyncio.wait(blocking_futures)
        # Lock the paths by a custom future.
        lock_future = asyncio.get_running_loop().create_future()
        for path in paths:
            assert path not in self._locks
            self._locks[path] = lock_future
            self._tracker.track(path)
        lock_future.add_done_callback(partial(self._locking_future_done_callback, paths))
        submitted = [False]
        try:

            def submitter(sub_op: Optional[AbstractIOQueueOperation] = None) -> asyncio.Future:
                if operation:
                    if sub_op:
                        raise RuntimeError("Do not provide an operation if the preparation has one.")
                    sub_op = operation
                elif not sub_op:
                    raise RuntimeError("Provide an operation at submit if the preparation has none.")
                elif sub_op.paths() != paths:
                    raise RuntimeError("Fs operation paths do not match prepared.")
                if submitted[0]:
                    raise RuntimeError("Fs operation already submitted.")
                submitted[0] = True
                sub_op.op_future = asyncio.get_running_loop().run_in_executor(self.pool, sub_op.execute)
                if not (ops := self._ops_by_account.get(sub_op.account_id)):
                    self._ops_by_account[sub_op.account_id] = ops = set()
                ops.add(sub_op)
                sub_op.op_future.add_done_callback(partial(self._operation_future_done_callback, sub_op, lock_future))
                return asyncio.shield(sub_op.op_future)

            yield submitter
        finally:
            # The submission callback was not used by the caller. The operation was not submitted, meaning the
            # blocking future must be cleaned up right now by making it completed.
            if not submitted[0]:
                lock_future.set_result(None)

    def _locking_future_done_callback(self, paths, f):
        for path in paths:
            assert self._locks[path] is f
            del self._locks[path]
            self._tracker.release(path)
        f.exception()

    def _operation_future_done_callback(
        self, operation: AbstractIOQueueOperation, locking_future: asyncio.Future, operation_future: asyncio.Future
    ):
        assert operation.op_future is operation_future
        ops = self._ops_by_account[operation.account_id]
        ops.remove(operation)
        if not ops:
            del self._ops_by_account[operation.account_id]
        del operation.op_future  # help cleanup
        if operation_future.cancelled():  # Running operations cannot be cancelled. See concurrent.Future
            locking_future.cancel()
        elif e := operation_future.exception():
            locking_future.set_exception(e)
        else:
            locking_future.set_result(operation_future.result())

    def _break_ex_locks(self, path: str) -> None:
        locks = []
        for p in self._token_tracker.affects(path):
            locks.append(self._token_locks[p])
        for lock in locks:
            self._del_ex_lock(lock)

    def _get_ex_lock_best(self, path: str) -> Optional[TokenLock]:
        any_locks = []
        old_locks = []
        now = time.time()
        try:
            for p in self._token_tracker.affects(path):
                lock = self._token_locks[p]
                if lock.until < now:
                    old_locks.append(lock)
                elif p == path:
                    return lock
                else:
                    any_locks.append(lock)
        finally:
            for old_lock in old_locks:
                self._del_ex_lock(old_lock)
        if any_locks:
            return min(any_locks, key=lambda l: (l.until, l.token))

    def _get_ex_lock_with_token(self, path: str, token: str) -> TokenLock:
        lock = self._get_ex_lock_best(path)
        if lock is None:
            raise TokenLockError("")
        elif lock.token != token:
            raise TokenLockError(lock.token)
        return lock

    def _add_ex_lock(self, lock: TokenLock) -> None:
        old_lock = self._token_locks.get(lock.path)
        assert old_lock is None
        self._token_locks[lock.path] = lock
        self._token_tracker.track(lock.path)

    def _del_ex_lock(self, lock: TokenLock) -> None:
        old_lock = self._token_locks.get(lock.path)
        if old_lock is lock:
            del self._token_locks[lock.path]
            self._token_tracker.release(lock.path)

    def get_ex_lock_token(self, path: str) -> Optional[str]:
        lock = self._get_ex_lock_best(path)
        if lock is not None:
            return lock.token
        for _ in self._tracker.affects(path):
            return ""  # locked but not by token

    def ex_lock(self, path: str, token: str) -> None:
        lock = self._get_ex_lock_best(path)
        if lock is None:
            for _ in self._tracker.affects(path):
                raise TokenLockError("")  # locked but not by token
            # create lock
            lock = TokenLock(path, token, time.time() + self.DEFAULT_TOKEN_LOCK_TTL)
            self._add_ex_lock(lock)
        elif lock.token == token:  # refresh lock
            lock.until = time.time() + self.DEFAULT_TOKEN_LOCK_TTL
        else:
            raise TokenLockError(lock.token)

    def ex_refresh_lock(self, path: str, token: str) -> None:
        lock = self._get_ex_lock_with_token(path, token)
        lock.until = time.time() + self.DEFAULT_TOKEN_LOCK_TTL

    def ex_unlock(self, path: str, token: str) -> None:
        lock = self._get_ex_lock_with_token(path, token)
        self._del_ex_lock(lock)

    def ex_replace_lock(self, path: str, token: str, old_token: str) -> None:
        self.ex_unlock(path, old_token)
        # NOTE: Since unlocking is guaranteed to unlock the exact lock if present, we'll surely be able to lock it.
        lock = TokenLock(path, token, time.time() + self.DEFAULT_TOKEN_LOCK_TTL)
        self._add_ex_lock(lock)
