import asyncio
import secrets
from abc import ABC, abstractmethod
from datetime import timedelta
from functools import wraps
from hashlib import sha256
from pathlib import Path
from time import time
from typing import Callable, Any, Union, Dict, Tuple, Optional
import pickle
from weakref import WeakValueDictionary

from aiofiles import os as async_os
from aiofiles import open as async_open

try:
    import orjson
    _use_orjson = True
except ImportError:
    import json
    _use_orjson = False

from vayu.log import L
from vayu.time import time_now, from_timestamp


class _MethodDecoratorAdaptor(object):
    def __init__(self, decorator, func):
        self.decorator = decorator
        self.func = func

    def __call__(self, *args, **kwargs):
        return self.decorator(self.func)(*args, **kwargs)

    def __get__(self, instance, owner):
        return self.decorator(self.func.__get__(instance, owner))


def auto_adapt_to_methods(decorator):
    """Allows you to use the same decorator on methods and functions,
    hiding the self argument from the decorator."""

    def adapt(func):
        return _MethodDecoratorAdaptor(decorator, func)

    return adapt


class Serializer(ABC):
    @abstractmethod
    def serialize(self, obj: Any) -> bytes:
        pass

    @abstractmethod
    def deserialize(self, bts: bytes) -> Any:
        pass


class Pickler(Serializer):
    def serialize(self, obj: Any) -> bytes:
        return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)

    def deserialize(self, bts: bytes) -> Any:
        return pickle.loads(bts)


class Jsoner(Serializer):
    def serialize(self, obj: Any) -> bytes:
        if _use_orjson:
            return orjson.dumps(obj)
        else:
            return json.dumps(obj).encode()

    def deserialize(self, bts: bytes) -> Any:
        if _use_orjson:
            return orjson.loads(bts)
        else:
            return json.loads(bts.decode())


class CacheMissError(Exception):
    pass


class Cache(ABC):
    def __init__(self, serializer: Serializer):
        self._serializer = serializer
        self._locks = dict()

    def cached(
            self,
            ttl: Union[timedelta, float],
            prefix: str = None,
            key: Union[str, Callable] = None,
            serializer: Serializer = None,
            log: bool = True,
    ):
        """Used as decorator to cache function calls.

        Args:
            ttl: cache expiry ttl
            prefix: If set, gets prepended to the rest of the cache key, otherwise uses function,
                        class, and module name to create prefix.
            key:
               - None: Use SHA256 digest of *args & **kwargs to construct cache key
               - str: Use the specified string and ignore function arguments
               - Callable: Calls it with *args * **kwargs which returns a str key.
            serializer: If not specified, default pickler serializer is used
            log: If True, logs reports such as cache hit, miss etc.

        """
        serializer = serializer or self._serializer
        if ttl is None:
            raise ValueError("ttl must be specified for cached decorator")

        if ttl is not None:
            if not isinstance(ttl, timedelta):
                ttl = timedelta(seconds=ttl)
            if ttl.total_seconds() <= 0:
                raise ValueError("ttl must be positive")

        @auto_adapt_to_methods
        def decorator(func):
            f_name = func.__name__
            if not asyncio.iscoroutinefunction(func):
                raise NotImplementedError("@cached can be used only with coroutines.")

            @wraps(func)
            async def async_inner(*args, **kwargs):
                t = time()
                k = self._key(func, prefix=prefix, key=key, args=args, kwargs=kwargs)
                t_del = time() - t
                if 0.001 < t_del and log:
                    L.warn(f"Key construction for cached fn `{f_name}` took too long: {t_del:.6f}s")

                try:
                    got = await self._read(k)
                    return serializer.deserialize(got)
                except CacheMissError:
                    pass

                async with self._locks.setdefault(k, asyncio.Lock()):
                    try:
                        got = await self._read(k)
                        return serializer.deserialize(got)
                    except CacheMissError:
                        pass

                t = time()
                result = await func(*args, **kwargs)
                t_del = time() - t
                if log:
                    L.info(f"@cached({f_name}) - CACHE MISS | took: {t_del:.6}s")
                await self._write(k, serializer.serialize(result), ttl)
                return result

            @wraps(func)
            async def evict(*args, **kwargs):
                k = self._key(func, prefix=prefix, key=key, args=args, kwargs=kwargs)
                await self._delete(k)

            async_inner.evict = evict

            return async_inner

        return decorator

    async def read(self, key: str) -> Any:
        got = await self._read(key)
        return self._serializer.deserialize(got)

    async def write(self, key: str, value: Any, ttl: timedelta):
        return await self._write(key, self._serializer.serialize(value), ttl)

    @abstractmethod
    async def _read(self, key: str) -> Any:
        raise NotImplementedError

    @abstractmethod
    async def _write(self, key: str, value: Any, ttl: timedelta):
        raise NotImplementedError

    @abstractmethod
    async def _delete(self, key: str):
        raise NotImplementedError

    @staticmethod
    def _key(
            fn: Callable,
            prefix: str = None,
            key: Union[str, Callable] = None,
            args: Tuple = None,
            kwargs: Dict = None,
    ) -> str:
        is_method = hasattr(fn, "__self__")

        if isinstance(key, str):
            return key

        parts = []
        if prefix:
            parts.append(prefix)
        else:
            parts.append(str(fn.__module__))
            if is_method:
                parts.append(str(fn.__self__.__class__.__name__))
            parts.append(fn.__name__)
        parts = [".".join(parts)]
        if callable(key):
            parts.append(key(*args, **kwargs))
        else:
            # Sort kwargs to ensure deterministic key.
            kwargs_key = [f"{k}={v}" for k, v in sorted(kwargs.items(), key=lambda x: x[0])]
            arg_key = ",".join([str(a) for a in args] + kwargs_key)
            parts.append(sha256(arg_key.encode()).hexdigest())
        return ":".join(parts)


class FileCache(Cache):
    def __init__(self, path: str, serializer: Serializer = None):
        super().__init__(serializer or Pickler())
        self._path = Path(path)
        self._ttl_bytes_length = 8
        self._path.mkdir(parents=True, exist_ok=True)

    def _filename(self, key: str) -> Path:
        return self._path / key

    async def _write(self, key: str, value: bytes, ttl: Optional[timedelta]):
        expiry_time = time_now(with_ms=True) + ttl
        ttl_val = int(expiry_time.timestamp() * 1e6)
        ttl_bytes = ttl_val.to_bytes(self._ttl_bytes_length, "big")
        bts = ttl_bytes + value

        final_path = self._filename(key)
        tmp_name = f".{final_path.name}.tmp.{secrets.token_hex(4)}"
        tmp_path = self._path / tmp_name

        async with async_open(tmp_path, "wb") as f:
            await f.write(bts)

        # Atomic replace (POSIX-atomic on same filesystem)
        await async_os.replace(tmp_path, final_path)

    async def _read(self, key: str) -> bytes:
        filename = self._filename(key)
        try:
            async with async_open(filename, "rb") as f:
                ttl_bytes = await f.read(self._ttl_bytes_length)
                if len(ttl_bytes) < self._ttl_bytes_length:
                    # corrupt file
                    raise CacheMissError()
                ttl_val = int.from_bytes(ttl_bytes, "big")
                expiry = from_timestamp(ttl_val / 1e6)  # reverse of write
                if expiry < time_now(with_ms=True):
                    raise CacheMissError()
                else:
                    bts = await f.read()
        except (FileNotFoundError, Exception):
            raise CacheMissError()

        return bts

    async def _delete(self, key: str):
        filename = self._filename(key)
        try:
            await async_os.remove(filename)
        except FileNotFoundError:
            pass


class MemoryCache(Cache):
    class NoOpSerializer(Serializer):
        def serialize(self, obj: Any) -> Any:
            return obj

        def deserialize(self, bts: Any) -> Any:
            return bts

    def __init__(self):
        super().__init__(self.NoOpSerializer())
        self._cache = {}

    async def _write(self, key: str, value: Any, ttl: timedelta):
        self._cache[key] = (time_now(with_ms=True) + ttl, value)

    async def _read(self, key: str) -> Optional[Any]:
        try:
            expiry_time, obj = self._cache[key]
        except KeyError:
            raise CacheMissError()
        if expiry_time < time_now(with_ms=True):
            await self._delete(key)
            raise CacheMissError()
        return obj

    async def _delete(self, key: str):
        try:
            del self._cache[key]
        except KeyError:
            L.info(f"Key {key} was not found in cache")
            pass


mem_cache = MemoryCache()
mem_cached = mem_cache.cached
