import asyncio
import secrets
import threading
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 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 __getattr__(self, item):
        return getattr(self.decorator(self.func), item)


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()
        self._sync_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 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):
            if asyncio.iscoroutinefunction(func):
                return self._make_async_wrapper(func, serializer, ttl, prefix, key, log)
            else:
                return self._make_sync_wrapper(func, serializer, ttl, prefix, key, log)

        return decorator

    def _make_async_wrapper(self, func, serializer, ttl, prefix, key, log):
        f_name = func.__name__

        @wraps(func)
        async def wrapper(*args, **kwargs):
            k = self._build_key_logged(func, prefix, key, args, kwargs, f_name, log)

            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)

        wrapper.evict = evict
        return wrapper

    def _make_sync_wrapper(self, func, serializer, ttl, prefix, key, log):
        f_name = func.__name__

        @wraps(func)
        def wrapper(*args, **kwargs):
            k = self._build_key_logged(func, prefix, key, args, kwargs, f_name, log)

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

            with self._sync_locks.setdefault(k, threading.Lock()):
                try:
                    got = self._read_sync(k)
                    return serializer.deserialize(got)
                except CacheMissError:
                    pass

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

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

        wrapper.evict = evict
        return wrapper

    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

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

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

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

    @staticmethod
    def _build_key_logged(func, prefix, key, args, kwargs, f_name, log) -> str:
        t = time()
        k = Cache._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")
        return k

    @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

    def _encode_entry(self, value: bytes, ttl: timedelta) -> bytes:
        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")
        return ttl_bytes + value

    def _decode_entry(self, data: bytes) -> bytes:
        if len(data) < self._ttl_bytes_length:
            raise CacheMissError()
        ttl_val = int.from_bytes(data[: self._ttl_bytes_length], "big")
        expiry = from_timestamp(ttl_val / 1e6)
        if expiry < time_now(with_ms=True):
            raise CacheMissError()
        return data[self._ttl_bytes_length :]

    async def _write(self, key: str, value: bytes, ttl: Optional[timedelta]):
        bts = self._encode_entry(value, ttl)

        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)

        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:
                data = await f.read()
        except (FileNotFoundError, Exception):
            raise CacheMissError()
        return self._decode_entry(data)

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

    def _write_sync(self, key: str, value: bytes, ttl: Optional[timedelta]):
        bts = self._encode_entry(value, ttl)

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

        with open(tmp_path, "wb") as f:
            f.write(bts)

        tmp_path.replace(final_path)

    def _read_sync(self, key: str) -> bytes:
        filename = self._filename(key)
        try:
            with open(filename, "rb") as f:
                data = f.read()
        except (FileNotFoundError, Exception):
            raise CacheMissError()
        return self._decode_entry(data)

    def _delete_sync(self, key: str):
        filename = self._filename(key)
        try:
            filename.unlink()
        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 = {}

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

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

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

    async def _write(self, key: str, value: Any, ttl: timedelta):
        self._write_sync(key, value, ttl)

    async def _read(self, key: str) -> Optional[Any]:
        return self._read_sync(key)

    async def _delete(self, key: str):
        self._delete_sync(key)


mem_cache = MemoryCache()
mem_cached = mem_cache.cached
