import asyncio
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.server.fstx.filesystem.operation.base import AbstractIOQueueOperation
from clarabot_nano.utility.path.tracker import PathDependencyTracker


class FilesystemTxQueue(object):
    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()

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

    @asynccontextmanager
    async def wait_ready(
        self, operation: Optional[AbstractIOQueueOperation] = None, paths: Optional[Sequence[str]] = None
    ) -> 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
            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())
