"""Bounded file work, isolated from authentication's default executor."""

from __future__ import annotations

import asyncio
import contextvars
import logging
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from weakref import WeakKeyDictionary

logger = logging.getLogger(__name__)


class FileIOPool:
    def __init__(self, workers: int, admitted: int, name: str):
        self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix=name)
        self._slots: WeakKeyDictionary = WeakKeyDictionary()
        self._admitted = admitted

    async def __call__(self, function, /, *args, **kwargs):
        return await self._run(function, args, kwargs, wait_on_cancel=False, on_cancel=None)

    async def run_until_complete(self, function, /, *args, on_cancel=None, **kwargs):
        """Keep a submitted worker alive until it finishes, even after cancellation.

        Callers can safely close inputs or remove temporary directories once this
        returns or raises. Cancellation while waiting for admission does not
        submit work. Once submitted, ``on_cancel`` may request cooperative
        cancellation; admission remains held until the actual worker exits.
        """
        return await self._run(function, args, kwargs, wait_on_cancel=True, on_cancel=on_cancel)

    async def _run(self, function, args, kwargs, *, wait_on_cancel, on_cancel):
        loop = asyncio.get_running_loop()
        semaphore = self._slots.setdefault(loop, asyncio.Semaphore(self._admitted))
        await semaphore.acquire()
        context = contextvars.copy_context()
        try:
            future = loop.run_in_executor(self._pool, context.run, partial(function, *args, **kwargs))
        except BaseException:
            semaphore.release()
            raise

        def finished(done):
            semaphore.release()
            if not done.cancelled():
                done.exception()

        future.add_done_callback(finished)
        # Cancellation cannot release admission before the real disk job ends.
        try:
            return await asyncio.shield(future)
        except asyncio.CancelledError:
            if not wait_on_cancel:
                raise
            if on_cancel is not None:
                try:
                    on_cancel()
                except Exception:
                    logger.exception("File worker cancellation callback failed")
            # The request owns the worker's inputs. Do not let its finally
            # block clean those up while a thread still reads or writes them.
            # Repeated Task.cancel() calls must not break that guarantee.
            while not future.done():
                try:
                    await asyncio.shield(future)
                except asyncio.CancelledError:
                    continue
                except BaseException:
                    break  # finished() consumes the worker's exception.
            raise


file_io = FileIOPool(4, 8, "artifact-files")
event_io = FileIOPool(2, 4, "event-spool")