# -*- coding: utf-8 -*-
"""
@File    : redis_manager.py
@Time    : 2023/9/7 18:42
@Author  : geekbing
@LastEditTime : -
@LastEditors : -
@Description : redis客户端Manager
"""
import asyncio
import functools
import inspect
import json
import pickle
from random import Random
from typing import Tuple

from awaits.awaitable import awaitable
from loguru import logger

from redis import ConnectionPool, StrictRedis
from rediscluster import RedisCluster, ClusterConnectionPool
from backend import settings

from apps.exceptions.error import RedisError


class RedisManager:
    """非线程安全，可能存在问题"""

    _cluster_pool = dict()
    _pool = dict()

    @property
    def client(self):
        pool = ConnectionPool(
            host=settings.REDIS_HOST,
            port=settings.REDIS_PORT,
            db=settings.REDIS_DB,
            max_connections=100,
            password=settings.REDIS_PASSWORD,
            decode_responses=True,
        )
        return StrictRedis(connection_pool=pool, decode_responses=True)

    @staticmethod
    def delete_client(redis_id: int, cluster: bool):
        """
        根据redis_id和是否是集群删除客户端
        :param redis_id:
        :param cluster:
        :return:
        """
        if cluster:
            RedisManager._cluster_pool.pop(redis_id)
        else:
            RedisManager._pool.pop(redis_id)

    @staticmethod
    def get_cluster_client(redis_id: int, address: str):
        """
        获取redis集群客户端
        :param redis_id:
        :param address:
        :return:
        """
        cluster = RedisManager._cluster_pool.get(redis_id)
        if cluster is not None:
            return cluster
        client = RedisManager.get_cluster(address)
        RedisManager._cluster_pool[redis_id] = client
        return client

    @staticmethod
    def get_single_node_client(redis_id: int, address: str, password: str, db: int):
        """
        获取redis单实例客户端
        :param redis_id:
        :param address:
        :param password:
        :param db:
        :return:
        """
        node = RedisManager._pool.get(redis_id)
        if node is not None:
            return node
        if ":" not in address:
            raise Exception("redis连接未包含端口号，请检查配置")
        host, port = address.split(":")
        pool = ConnectionPool(
            host=host,
            port=port,
            db=db,
            max_connections=100,
            password=password,
            decode_responses=True,
        )
        client = StrictRedis(connection_pool=pool)
        RedisManager._pool[redis_id] = client
        return client

    @staticmethod
    def refresh_redis_client(redis_id: int, address: str, password: str, db: str):
        """
        刷新redis客户端
        :param redis_id:
        :param address:
        :param password:
        :param db:
        :return:
        """
        host, port = address.split(":")
        pool = ConnectionPool(
            host=host,
            port=port,
            db=db,
            max_connections=100,
            password=password,
            decode_responses=True,
        )
        client = StrictRedis(connection_pool=pool, decode_responses=True)
        RedisManager._pool[redis_id] = client

    @staticmethod
    def refresh_redis_cluster(redis_id: int, addr: str):
        RedisManager._cluster_pool[redis_id] = RedisManager.get_cluster(addr)

    @staticmethod
    def get_cluster(address: str):
        """
        获取集群连接池
        :param address:
        :return:
        """
        try:
            nodes = address.split(",")
            startup_nodes = [
                {"host": n.split(":")[0], "port": n.split(":")[1]}
                for n in nodes
                if ":" in n
            ]
            if len(startup_nodes) == 0:
                raise Exception("找不到集群节点，请检查配置")
            pool = ClusterConnectionPool(
                startup_nodes=startup_nodes, max_connections=100, decode_responses=True
            )
            client = RedisCluster(connection_pool=pool, decode_responses=True)
            return client
        except Exception as e:
            raise RedisError(f"获取Redis连接失败, {e}")


class RedisHelper:
    prefix = "fastapi"
    redis_client = RedisManager().client

    @staticmethod
    @awaitable
    def execute_command(client, command, *args, **kwargs):
        return client.execute_command(command, *args, **kwargs)

    @staticmethod
    @awaitable
    def ping():
        """
        test redis client
        :return:
        """
        return RedisHelper.redis_client.ping()

    @staticmethod
    def get_address_record(address: str):
        """
        获取ip是否已经开启录制

        :param address:
        :return:
        """
        key = RedisHelper.get_key(f"record:ip:{address}")
        return RedisHelper.redis_client.get(key)

    @staticmethod
    def get_user_record(user_id: str):
        """
        获取当前用户是否已开启录制

        :param user_id:
        :return:
        """
        key = RedisHelper.get_key(f"user:id:{user_id}")
        return RedisHelper.redis_client.get(key)

    @staticmethod
    @awaitable
    def cache_record(user_id: str, request):
        """
        缓存录制数据

        :param user_id: 开启录制的用户id
        :param request: 客户端请求流量
        :return:
        """
        key = RedisHelper.get_key(f"id:{user_id}:requests")
        RedisHelper.redis_client.rpush(key, request)
        ttl = RedisHelper.redis_client.ttl(key)
        if ttl < 0:
            RedisHelper.redis_client.expire(key, 3600)

    @staticmethod
    def set_address_record(
        user_id: int,
        address: str,
        regex: str,
        is_local: bool,
    ):
        """
        设置录制状态

        :param user_id:
        :param address:
        :param regex: 录制的url正则
        :param is_local: False: 其他端录制，True: 本机录制
        :return:
        """
        # 默认录制数据保存1小时
        value = json.dumps(
            {
                "user_id": user_id,
                "regex": regex,
                "ip": address,
                "local": is_local,
            },
            ensure_ascii=False,
        )
        RedisHelper.redis_client.set(
            RedisHelper.get_key(f"record:ip:{address}"), value, ex=3600
        )
        RedisHelper.redis_client.set(
            RedisHelper.get_key(f"user:id:{user_id}"), value, ex=3600
        )
        # 清除上次录制数据
        RedisHelper.redis_client.delete(RedisHelper.get_key(f"id:{user_id}:requests"))

    @staticmethod
    def remove_address_record(address: str):
        """
        停止录制任务

        :param address:
        :return:
        """
        return RedisHelper.redis_client.delete(
            RedisHelper.get_key(f"record:ip:{address}")
        )

    @staticmethod
    def remove_user_record(user_id: str):
        """
        停止录制任务

        :param user_id:
        :return:
        """
        return RedisHelper.redis_client.delete(
            RedisHelper.get_key(f"user:id:{user_id}")
        )

    @staticmethod
    def list_record_data(user_id: str):
        """
        获取录制数据

        :param user_id: 开启录制的用户id
        :return:
        """
        key = RedisHelper.get_key(f"id:{user_id}:requests")
        data = RedisHelper.redis_client.lrange(key, 0, -1)
        return [json.loads(x) for x in data]

    @staticmethod
    def remove_record_data(user_id: str, index: int):
        """
        删除录制数据

        :param user_id: 开启录制的用户id
        :param index:
        :return:
        """
        key = RedisHelper.get_key(f"id:{user_id}:requests")
        RedisHelper.redis_client.lset(key, index, "DELETED")
        RedisHelper.redis_client.lrem(key, 1, "DELETED")

    @staticmethod
    def async_delete_prefix(key: str):
        """
        根据前缀删除数据
        :param key:
        :return:
        """
        for k in RedisHelper.redis_client.scan_iter(f"{key}*"):
            RedisHelper.redis_client.delete(k)
            logger.bind(name=None).debug(f"delete redis key: {k}")

    @staticmethod
    def delete_prefix(key: str):
        """
        根据前缀删除数据
        :param key:
        :return:
        """
        for k in RedisHelper.redis_client.scan_iter(f"{key}*"):
            RedisHelper.redis_client.delete(k)
            logger.bind(name=None).debug(f"delete redis key: {k}")

    @staticmethod
    def get_key(_redis_key: str, args_key: bool = True, *args, **kwargs):
        if not args_key:
            return f"{RedisHelper.prefix}:{_redis_key}"
        filter_args = [
            a
            for a in args
            if not str(a).startswith(("<class", "<sqlalchemy", "(<sqlalchemy"))
        ]
        for v in kwargs.values():
            if v and not str(v).startswith(("<class", "<sqlalchemy", "(<sqlalchemy")):
                filter_args.append(str(v))
        return (
            f"{RedisHelper.prefix}:{_redis_key}"
            f"{':' + ':'.join(str(a) for a in filter_args) if len(filter_args) > 0 else ''}"
        )

    @staticmethod
    def get_key_with_suffix(cls_name: str, key: str, args: tuple, key_suffix):
        filter_args = [a for a in args if not str(args[0]).startswith("<class")]
        suffix = key_suffix(filter_args)
        return f"{RedisHelper.prefix}:{cls_name}:{key}:{suffix}"

    @staticmethod
    def cache(key: str, expired_time=30 * 60, args_key=True):
        """
        自动缓存装饰器
        :param args_key:
        :param key: 被缓存的key
        :param expired_time: 默认key过期时间
        :return:
        """

        def decorator(func):
            # 缓存已存在
            if asyncio.iscoroutinefunction(func):

                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                    # TODO: redis重复代码优化
                    if not settings.REDIS_ON:
                        return func(*args, **kwargs)
                    cls_name = (
                        inspect.getframeinfo(inspect.currentframe().f_back)[3][0]
                        .split(".")[0]
                        .split(" ")[-1]
                    )
                    redis_key = RedisHelper.get_key(
                        f"{cls_name}:{key}", args_key, *args, **kwargs
                    )
                    data = RedisHelper.redis_client.get(redis_key)
                    # 缓存已存在
                    if data is not None:
                        return pickle.loads(bytes.fromhex(data))
                    # 获取最新数据
                    new_data = func(*args, **kwargs)
                    info = pickle.dumps(new_data)
                    # logger.bind(name=None).debug(f"set redis key: {redis_key}")
                    RedisHelper.redis_client.set(redis_key, info.hex(), ex=expired_time)
                    return new_data

                return wrapper
            else:

                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                    if not settings.REDIS_ON:
                        return func(*args, **kwargs)
                    cls_name = (
                        inspect.getframeinfo(inspect.currentframe().f_back)[3][0]
                        .split(".")[0]
                        .split(" ")[-1]
                    )
                    redis_key = RedisHelper.get_key(
                        f"{cls_name}:{key}", args_key, *args, **kwargs
                    )
                    data = RedisHelper.redis_client.get(redis_key)
                    # 缓存已存在
                    if data is not None:
                        return pickle.loads(bytes.fromhex(data))
                    # 获取最新数据
                    new_data = func(*args, **kwargs)
                    info = pickle.dumps(new_data)
                    # logger.bind(name=None).debug(f"set redis key: {redis_key}")
                    # 添加随机数防止缓存雪崩
                    RedisHelper.redis_client.set(
                        redis_key,
                        info.hex(),
                        ex=expired_time + Random().randint(10, 59),
                    )
                    return new_data

                return wrapper

        return decorator

    @staticmethod
    def up_cache(*key: str, key_and_suffix: Tuple = None):
        """
        redis缓存key，套了此方法，会自动执行更新数据操作后删除缓存
        :param key:
        :param key_and_suffix: 要删除的key和key组成规则
        :return:
        """

        def decorator(func):
            if asyncio.iscoroutinefunction(func):

                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                    new_data = func(*args, **kwargs)
                    if not settings.REDIS_ON:
                        return new_data
                    cls_name = (
                        inspect.getframeinfo(inspect.currentframe().f_back)[3][0]
                        .split(".")[0]
                        .split(" ")[-1]
                    )
                    for k in key:
                        redis_key = f"{RedisHelper.prefix}:{cls_name}:{k}"
                        RedisHelper.async_delete_prefix(redis_key)
                    if key_and_suffix is not None:
                        current_key = RedisHelper.get_key_with_suffix(
                            cls_name, key_and_suffix[0], args, key_and_suffix[1]
                        )
                        RedisHelper.redis_client.delete(current_key)
                    # 更新数据，删除缓存
                    return new_data

                return wrapper
            else:

                @functools.wraps(func)
                def wrapper(*args, **kwargs):
                    new_data = func(*args, **kwargs)
                    if not settings.REDIS_ON:
                        return new_data
                    cls_name = (
                        inspect.getframeinfo(inspect.currentframe().f_back)[3][0]
                        .split(".")[0]
                        .split(" ")[-1]
                    )
                    for k in key:
                        redis_key = f"{RedisHelper.prefix}:{cls_name}:{k}"
                        RedisHelper.delete_prefix(redis_key)
                    if key_and_suffix is not None:
                        current_key = RedisHelper.get_key_with_suffix(
                            cls_name, key_and_suffix[0], args, key_and_suffix[1]
                        )
                        RedisHelper.redis_client.delete(current_key)
                    return new_data

                return wrapper

        return decorator