import abc
import logging
import os
from abc import abstractmethod
from collections import defaultdict
from inspect import isasyncgenfunction, iscoroutinefunction, isgeneratorfunction
from typing import Callable, DefaultDict, Optional, Type  # ParamSpec

from uhura.composition import async_unit, compose
from uhura.serde import DEFAULT_SERDE, Serde
from uhura.base import Cacheable

logger = logging.getLogger("uhura.caches")


class Cache(abc.ABC):
    @abstractmethod
    def exists(self) -> bool:
        raise NotImplementedError()

    @abstractmethod
    def get(self):
        raise NotImplementedError()

    @abstractmethod
    def update(self, obj):
        raise NotImplementedError()


class LocalCache(Cache):
    def __init__(self, path: str, serde=DEFAULT_SERDE) -> None:
        self._path = path
        self._serde = serde

    def exists(self):
        return os.path.exists(self._path)

    def update(self, obj):
        logger.info(f"Dumping to {self._path}")

        if not os.path.exists(os.path.dirname(self._path)):
            os.makedirs(os.path.dirname(self._path), exist_ok=True)
        self._serde.write_to_file(self._path, obj)
        return obj

    def get(self):
        return self._serde.read_from_file(self._path)

    @classmethod
    def render_path(cls, base_path: str, cache_key: str, serde: Serde):
        if not cache_key.endswith(serde.file_extension):
            cache_key = cache_key + serde.file_extension
        return os.path.join(base_path, cache_key)

    @classmethod
    def cache_for_instance(cls, cacheable: Cacheable, base_path: str):
        cache_key = cacheable.cache_key()
        serde = cacheable.get_serde()
        path = cls.render_path(base_path, cache_key, serde)
        return cls(path, serde)


async def _take(async_iterable, n=5):
    index = 0
    async for item in async_iterable:
        index += 1
        if index > n:
            break
        yield item


def _sync_take(iterable, n=5):
    index = 0
    for item in iterable:
        index += 1
        if index > n:
            break
        yield item


class CachingService:
    def __init__(self, lookup, base_path, depth=4):
        self._base_path = base_path
        self._depth = depth
        self._lookup: DefaultDict[str, Callable[[Cacheable, str, ...], Cache]] = lookup

    def _iterate_from_cache(self, instance: Cacheable):
        for _ in range(self._depth):
            cache = self.cache_for_instance(instance)
            yield cache.get()
            instance._read_count += 1

    def cache_for_instance(self, instance: Cacheable):
        return self._lookup[instance.cache_key()](instance, self._base_path)

    def store_and_return(self, instance: Cacheable, method):
        def _update(value):
            cache = self.cache_for_instance(instance)
            if cache.exists():
                return cache.get()
            else:
                cache.update(value)
                return value

        if isasyncgenfunction(method):
            return _take(compose(method, _update)(instance), n=self._depth)
        if isgeneratorfunction(method):
            return _sync_take(compose(method, _update)(instance), n=self._depth)
        return compose(method, _update)(instance)

    def pull_from_cache(self, instance: Cacheable, original_method):
        if isasyncgenfunction(original_method):
            return async_unit(self._iterate_from_cache(instance))
        if isgeneratorfunction(original_method):
            return self._iterate_from_cache(instance)
        if iscoroutinefunction(original_method):
            return async_unit(self.cache_for_instance(instance).get())
        return self.cache_for_instance(instance).get()

    @classmethod
    def default(cls, base_path, **kwargs):
        return cls(
            lookup=defaultdict(lambda: LocalCache.cache_for_instance),
            base_path=base_path,
            **kwargs,
        )
