import functools
import os
from dataclasses import dataclass
from typing import Any, Callable, Optional, TypeVar, ParamSpec, cast

import numpy as np
from diskcache import Cache, Timeout
from prettytable import PrettyTable

from zyplib._config import config
from zyplib.utils.print import print_info, print_warn, type_summary

_SENTINEL = object()

_configs = {
    'cache_dir': config.DISK_CACHE_DIR,
    'cache_max_size': config.DISK_CACHE_MAX_SIZE,  # MB
}


def update_cache_config(
    cache_dir: str | None = None, cache_max_size_mb: int | float | None = None
):
    if cache_dir is not None:
        _configs['cache_dir'] = cache_dir
    if cache_max_size_mb is not None:
        _configs['cache_max_size'] = cache_max_size_mb


def cache_factory(cache_dir: str | None = None) -> Cache:
    """创建缓存对象

    - 如果传入 `cache_dir` 则使用传入的目录作为 diskcache 的缓存目录
    - 如果传入 `cache_dir` 为 None 则使用默认缓存目录
        - 默认缓存目录为 `./cache`, 可以通过 `update_cache_config` 函数设置

    Parameters
    ----------
    - `cache_dir` : `str`, optional
        - 缓存目录; 如果为 None 则使用 _config 中的默认缓存目录

    Returns
    ----------
    - `Cache`
        - 缓存对象
    """
    cache_dir = cache_dir or _configs['cache_dir']
    if not os.path.exists(cache_dir):
        os.makedirs(cache_dir)
    size_limit = int(_configs['cache_max_size'] * 1024 * 1024)
    return Cache(cache_dir, size_limit=size_limit)


def get_cache(key: str, default: Any | None = None, cache_dir: str | None = None):
    """获取缓存

    Parameters
    ----------
    - `key` : `str`
        - 缓存键
    - `default` : `Any`, optional
        - 默认值，当缓存不存在时返回
    - `cache_dir` : `str | None`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录

    Returns
    ----------
    - `Any`
        - 缓存值, 当缓存不存在或者过期时返回默认值
    """
    try:
        with cache_factory(cache_dir) as cache:
            return cache.get(key, default)
    except Timeout as t:
        print_warn(f'Cache[{key}] timeout: {t}')
        return default


def set_cache(
    key: str,
    value: Any,
    expire_seconds: int | None = None,
    tag: str | None = None,
    cache_dir: str | None = None,
):
    """设置缓存

    Parameters
    ----------
    - `key` : `str`
        - 缓存键
    - `value` : `Any`
        - 缓存值
    - `expire_seconds` : `int`, optional
        - 缓存过期时间; 如果为 None 则不设置过期时间
    - `tag` : `str`, optional
        - 缓存标签
    - `cache_dir` : `str | None`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录

    Returns
    ----------
    - `bool`
        - 设置成功返回 True, 否则返回 False
    """
    with cache_factory(cache_dir) as cache:
        flag = cache.set(key, value, expire=expire_seconds, tag=tag)
    if not flag:
        print_warn(f'Cache[{key}] set failed')
    return flag


def del_cache(key: str, cache_dir: str | None = None):
    """删除缓存

    Parameters
    ----------
    - `key` : `str`
        - 缓存键
    - `cache_dir` : `str | None`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录
    """
    with cache_factory(cache_dir) as cache:
        cache.delete(key)


def purge_all_cache(cache_dir: str | None = None):
    """清空所有缓存

    Parameters
    ----------
    - `cache_dir` : `str | None`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录
    """
    with cache_factory(cache_dir) as cache:
        cache.clear()


def use_cache(
    key,
    fn_value: Callable[[], Any],
    cache_dir: str | None = None,
    expire_seconds: int | None = None,
    tag: str | None = None,
):
    """使用缓存, 如果缓存命中则直接返回缓存值, 否则调用 fn_value 计算值并缓存


    Parameters
    ----------
    - `key` : `str`
        - 缓存键
    - `fn_value` : `Callable[[], Any]`
        - 计算值的函数
    - `cache_dir` : `str | None`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录
    - `expire_seconds` : `int`, optional
        - 缓存过期时间; 如果为 None 则不设置过期时间
    - `tag` : `str`, optional
        - 缓存标签

    Returns
    ----------
    - `Any`
        - 缓存值

    Examples
    ----------
    ```py
    def fn_value():
        return 1
    use_cache('key', fn_value)
    ```
    """
    with cache_factory(cache_dir) as cache:
        value = cache.get(key, default=_SENTINEL)
        if value is _SENTINEL:
            value = fn_value()
            cache.set(key, value, expire=expire_seconds, tag=tag)
    return value


@dataclass
class CacheItem:
    key: Any
    value: Any
    expire_time: float | None
    tag: Any | None


def list_caches(cache_dir: str | None = None):
    """列出所有缓存; 通过

    Parameters
    ----------
    - `cache_dir` : `str`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录

    Yields
    ----------
    - `CacheItem`: 缓存项
    """
    with cache_factory(cache_dir) as cache:
        for key in cache.iterkeys():
            entry = cache.get(key, default=None, expire_time=True, tag=True)
            if not isinstance(entry, tuple) or len(entry) != 3:
                raise TypeError(
                    f'Invalid cache entry structure for key={key!r}: {type(entry).__name__}'
                )

            obj, expire_time, tag = cast(tuple[Any, Any, Any], entry)
            cache_item = CacheItem(
                key=key,
                value=obj,
                expire_time=None if expire_time is None else float(expire_time),
                tag=tag,
            )
            yield cache_item


def display_caches(cache_dir: str | None = None, table: bool = True):
    """打印缓存状态

    Parameters
    ----------
    - `cache_dir` : `str`, optional
        - 缓存目录; 如果为 None 则使用 cache_factory 中的默认缓存目录
    - `table` : `bool`, optional
        - 是否使用 PrettyTable 展示; 如果为 False 则使用 print 打印
    """
    with cache_factory(cache_dir) as cache:
        size = cache.volume()
    print_info(f'缓存大小: {size / 1024 / 1024:.2f} MB')

    caches = list_caches(cache_dir)
    if table:
        table_obj = PrettyTable(['Index', 'Key', 'Value', 'Expire', 'Tag'])
        for idx, cache_item in enumerate(caches):
            table_obj.add_row(
                [
                    idx,
                    type_summary(cache_item.key),
                    type_summary(cache_item.value),
                    cache_item.expire_time,
                    cache_item.tag,
                ]
            )
        print_info(f'缓存: {cache_dir or _configs["cache_dir"]}'.center(60, '='))
        print_info(table_obj)
    else:
        print_info(f'缓存: {cache_dir or _configs["cache_dir"]}')
        for idx, cache_item in enumerate(caches, start=1):
            print_info(f'缓存 {idx}'.center(60, '='))
            print_info('Key:', end='\t')
            print(cache_item.key)
            print_info('Value:', end='\t')
            print(cache_item.value)
            print_info('Expire:', end='\t')
            print(cache_item.expire_time)
            print_info('Tag:', end='\t')
            print(cache_item.tag)
            print('')


P = ParamSpec('P')
T = TypeVar('T')


def make_func_memoize(
    func: Callable[P, T],
    cache: Cache | None = None,
    name=None,
    typed=False,
    expire=None,
    tag=None,
    ignore=(),
) -> Callable[P, T]:
    """接受一个函数对象，为其动态添加 diskcache 的 memoize 功能

    Parameters
    ----------
    - `func` : `Callable`
        - 需要添加缓存功能的函数
    - `cache` : `Cache`
        - diskcache 的 Cache 对象
    - `name` : `str`, optional
        - 缓存的名称; 如果为 None 则使用函数名
    - `typed` : `bool`, optional
        - 是否区分参数类型; 如果为 True 则不同类型的参数会被
        - 缓存不同的结果
    - `expire` : `int`, optional
        - 缓存过期时间 seconds; 如果为 None 则不设置过期时间
    - `tag` : `str`, optional
        - 缓存标签; 如果为 None 则不设置标签

    Examples
    ----------
    ```py
    @make_func_memoize  # 不能有参数
    def fn_value():
        return 1


    f1_cached = make_func_memoize(f1, expire=60*60)
    ```
    """
    if cache is None:
        cache = cache_factory()

    memoized_func = cache.memoize(
        name=name, typed=typed, expire=expire, tag=tag, ignore=ignore
    )(func)

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        return cast(T, memoized_func(*args, **kwargs))

    return wrapper
