import asyncio
from concurrent.futures import Future as ConcurrentFuture
from typing import Any, Callable, Generic, TypeVar

from ..base import BaseModel

T = TypeVar("T")
U = TypeVar("U")
M = TypeVar("M", bound=BaseModel)


class APIFuture(Generic[T]):
    def __init__(self, future: ConcurrentFuture[T]):
        self._future = future

    @staticmethod
    def completed(value=None):
        f = ConcurrentFuture()
        f.set_result(value)
        return APIFuture(f)

    def result(self, timeout: float | None = None) -> T:
        return self._future.result(timeout=timeout)

    def done(self) -> bool:
        return self._future.done()

    def exception(self, timeout: float | None = None):
        return self._future.exception(timeout=timeout)

    def cancel(self) -> bool:
        return self._future.cancel()

    def map(self, fn: Callable[[T], U]) -> "APIFuture[U]":
        new_future: ConcurrentFuture = ConcurrentFuture()

        def _callback(src: ConcurrentFuture):
            try:
                value = src.result()
                new_value = fn(value)
                new_future.set_result(new_value)
            except Exception as e:
                new_future.set_exception(e)

        self._future.add_done_callback(_callback)
        return APIFuture(new_future)

    def map_model(
        self: "APIFuture[dict[str, Any]]", model_cls: type[M]
    ) -> "APIFuture[M]":
        return self.map(lambda r: model_cls.model_validate(r.get("result", {})))

    def map_dict(self: "APIFuture[dict[str, Any]]") -> "APIFuture[dict]":
        return self.map(lambda r: r.get("result", {}))

    def __await__(self):
        wrapped = asyncio.wrap_future(self._future)
        return wrapped.__await__()
