from __future__ import annotations

import asyncio
import dataclasses
from abc import ABC, abstractmethod
from asyncio import create_task, gather, get_event_loop
from asyncio.futures import Future
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Generic,
    TypeVar,
    overload,
)

from .exceptions import WrongNumberOfResultsReturned

if TYPE_CHECKING:
    from asyncio.events import AbstractEventLoop
    from collections.abc import (
        Awaitable,
        Callable,
        Hashable,
        Iterable,
        Mapping,
        Sequence,
    )


T = TypeVar("T")
K = TypeVar("K")


@dataclass
class LoaderTask(Generic[K, T]):
    key: K
    future: Future


@dataclass
class Batch(Generic[K, T]):
    tasks: list[LoaderTask] = dataclasses.field(default_factory=list)
    dispatched: bool = False
    _dispatch_task: asyncio.Task[None] | None = dataclasses.field(
        default=None, repr=False
    )

    def add_task(self, key: Any, future: Future) -> None:
        task = LoaderTask[K, T](key, future)
        self.tasks.append(task)
        future.add_done_callback(self._on_future_done)

    def _on_future_done(self, future: Future) -> None:
        """Cancel the dispatch task if all futures in the batch are cancelled."""
        if (
            future.cancelled()
            and self._dispatch_task is not None
            and not self._dispatch_task.done()
            and all(t.future.cancelled() for t in self.tasks)
        ):
            self._dispatch_task.cancel()

    def __len__(self) -> int:
        return len(self.tasks)


class AbstractCache(ABC, Generic[K, T]):
    @abstractmethod
    def get(self, key: K) -> Future[T] | None:
        pass

    @abstractmethod
    def set(self, key: K, value: Future[T]) -> None:
        pass

    @abstractmethod
    def delete(self, key: K) -> None:
        pass

    @abstractmethod
    def clear(self) -> None:
        pass


class DefaultCache(AbstractCache[K, T]):
    def __init__(self, cache_key_fn: Callable[[K], Hashable] | None = None) -> None:
        self.cache_key_fn: Callable[[K], Hashable] = (
            cache_key_fn if cache_key_fn is not None else lambda x: x
        )
        self.cache_map: dict[Hashable, Future[T]] = {}

    def get(self, key: K) -> Future[T] | None:
        return self.cache_map.get(self.cache_key_fn(key))

    def set(self, key: K, value: Future[T]) -> None:
        self.cache_map[self.cache_key_fn(key)] = value

    def delete(self, key: K) -> None:
        del self.cache_map[self.cache_key_fn(key)]

    def clear(self) -> None:
        self.cache_map.clear()


class DataLoader(Generic[K, T]):
    batch: Batch[K, T] | None = None
    cache: bool = False
    cache_map: AbstractCache[K, T]

    @overload
    def __init__(
        self,
        # any BaseException is rethrown in 'load', so should be excluded from the T type
        load_fn: Callable[[list[K]], Awaitable[Sequence[T | BaseException]]],
        max_batch_size: int | None = None,
        cache: bool = True,
        loop: AbstractEventLoop | None = None,
        cache_map: AbstractCache[K, T] | None = None,
        cache_key_fn: Callable[[K], Hashable] | None = None,
    ) -> None: ...

    # fallback if load_fn is untyped and there's no other info for inference
    @overload
    def __init__(
        self: DataLoader[K, Any],
        load_fn: Callable[[list[K]], Awaitable[list[Any]]],
        max_batch_size: int | None = None,
        cache: bool = True,
        loop: AbstractEventLoop | None = None,
        cache_map: AbstractCache[K, T] | None = None,
        cache_key_fn: Callable[[K], Hashable] | None = None,
    ) -> None: ...

    def __init__(
        self,
        load_fn: Callable[[list[K]], Awaitable[Sequence[T | BaseException]]],
        max_batch_size: int | None = None,
        cache: bool = True,
        loop: AbstractEventLoop | None = None,
        cache_map: AbstractCache[K, T] | None = None,
        cache_key_fn: Callable[[K], Hashable] | None = None,
    ):
        self.load_fn = load_fn
        self.max_batch_size = max_batch_size

        self._loop = loop

        self.cache = cache

        if self.cache:
            self.cache_map = (
                DefaultCache(cache_key_fn) if cache_map is None else cache_map
            )

    @property
    def loop(self) -> AbstractEventLoop:
        if self._loop is None:
            self._loop = get_event_loop()

        return self._loop

    def load(self, key: K) -> Awaitable[T]:
        if self.cache:
            future = self.cache_map.get(key)

            if future and not future.cancelled():
                return future

        future = self.loop.create_future()

        if self.cache:
            self.cache_map.set(key, future)

        batch = get_current_batch(self)
        batch.add_task(key, future)

        return future

    def load_many(self, keys: Iterable[K]) -> Awaitable[list[T]]:
        return gather(*map(self.load, keys))

    def clear(self, key: K) -> None:
        if self.cache:
            self.cache_map.delete(key)

    def clear_many(self, keys: Iterable[K]) -> None:
        if self.cache:
            for key in keys:
                self.cache_map.delete(key)

    def clear_all(self) -> None:
        if self.cache:
            self.cache_map.clear()

    def prime(self, key: K, value: T, force: bool = False) -> None:
        self.prime_many({key: value}, force)

    def prime_many(self, data: Mapping[K, T], force: bool = False) -> None:
        # Populate the cache with the specified values
        if self.cache:
            for key, value in data.items():
                if not self.cache_map.get(key) or force:
                    future: Future = Future(loop=self.loop)
                    future.set_result(value)
                    self.cache_map.set(key, future)

        # For keys that are pending on the current batch, but the
        # batch hasn't started fetching yet: Remove it from the
        # batch and set to the specified value
        if self.batch is not None and not self.batch.dispatched:
            batch_updated = False
            for task in self.batch.tasks:
                if task.key in data:
                    batch_updated = True
                    task.future.set_result(data[task.key])
            if batch_updated:
                self.batch.tasks = [
                    task for task in self.batch.tasks if not task.future.done()
                ]


def should_create_new_batch(loader: DataLoader, batch: Batch) -> bool:
    return bool(
        batch.dispatched
        or (loader.max_batch_size and len(batch) >= loader.max_batch_size)
    )


def get_current_batch(loader: DataLoader) -> Batch:
    if loader.batch and not should_create_new_batch(loader, loader.batch):
        return loader.batch

    loader.batch = Batch()

    dispatch(loader, loader.batch)

    return loader.batch


def dispatch(loader: DataLoader, batch: Batch) -> None:
    def _schedule() -> None:
        batch._dispatch_task = create_task(dispatch_batch(loader, batch))
        batch._dispatch_task.add_done_callback(
            lambda _: setattr(batch, "_dispatch_task", None)
        )

    loader.loop.call_soon(_schedule)


async def dispatch_batch(loader: DataLoader, batch: Batch) -> None:
    batch.dispatched = True

    keys = [task.key for task in batch.tasks]
    if len(keys) == 0:
        # Ensure batch is not empty
        # Unlikely, but could happen if the tasks are
        # overridden with preset values
        return

    # Skip the batch entirely if all futures have already been cancelled.
    # This avoids calling load_fn (e.g. a database query) when no caller
    # is waiting for the results.
    if all(task.future.cancelled() for task in batch.tasks):
        return

    # TODO: check if load_fn return an awaitable and it is a list

    try:
        values = await loader.load_fn(keys)
        values = list(values)

        if len(values) != len(batch):
            raise WrongNumberOfResultsReturned(  # noqa: TRY301
                expected=len(batch), received=len(values)
            )

        for task, value in zip(batch.tasks, values, strict=True):
            # Trying to set_result in a cancelled future would raise
            # asyncio.exceptions.InvalidStateError
            if task.future.cancelled():
                continue

            if isinstance(value, BaseException):
                task.future.set_exception(value)
            else:
                task.future.set_result(value)
    except Exception as e:  # noqa: BLE001
        for task in batch.tasks:
            if task.future.cancelled():
                continue
            task.future.set_exception(e)


__all__ = [
    "AbstractCache",
    "Batch",
    "DataLoader",
    "DefaultCache",
    "LoaderTask",
    "dispatch",
    "dispatch_batch",
    "get_current_batch",
    "should_create_new_batch",
]
