"""
@File    :   redis_client.py
@Time    :   2016/05/10
@Author  :   lijc210@163.com
@Desc    :
Redis 连接客户端，支持单机模式和集群模式
"""

from typing import Any, Dict, List, Optional, Tuple, Union

import redis
from redis.cluster import RedisCluster


class RedisClient:
    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        db: Optional[int] = None,
        password: Optional[str] = None,
        pool: bool = False,
        startup_nodes: Optional[List[Dict[str, Any]]] = None,
        decode_responses: bool = True,
    ) -> None:
        """
        Redis 连接客户端

        :param host: Redis 服务器地址
        :param port: Redis 服务器端口
        :param db: 数据库编号（单机模式）
        :param password: 密码
        :param pool: 是否使用连接池
        :param startup_nodes: 集群节点列表，如 [{"host": "x.x.x.x", "port": 7000}]
        :param decode_responses: 是否解码响应，默认 True
        """
        self.host = host
        self.port = port
        self.db = db
        self.password = password
        self.pool = pool
        self.startup_nodes = startup_nodes
        self.decode_responses = decode_responses
        self.client = self.conn()

    def conn(self) -> Union[redis.Redis, RedisCluster]:
        """
        连接 Redis
        集群模式只有一个 db

        :return: Redis 客户端对象
        """
        if self.startup_nodes:
            # 连接集群
            conn = RedisCluster(
                startup_nodes=self.startup_nodes,
                decode_responses=self.decode_responses,
                password=self.password,
            )
        elif self.pool:
            # 连接池模式
            pool = redis.ConnectionPool(
                host=self.host,
                port=self.port,
                db=self.db,
                password=self.password,
                decode_responses=self.decode_responses,
            )
            conn = redis.Redis(connection_pool=pool)
        else:
            # 单机模式
            conn = redis.StrictRedis(
                host=self.host,
                port=self.port,
                db=self.db,
                password=self.password,
                decode_responses=self.decode_responses,
            )
        return conn

    # ============ String 操作 ============

    def set(self, key: str, value: Any, ex: Optional[int] = None) -> bool:
        """
        设置字符串类型

        :param key: 键
        :param value: 值
        :param ex: 过期时间（秒）
        :return: 是否成功
        """
        if ex:
            return self.client.set(key, value, ex=ex)
        return self.client.set(key, value)

    def get(self, key: str) -> Optional[Any]:
        """
        获取字符串值

        :param key: 键
        :return: 值
        """
        return self.client.get(key)

    def mset(self, mapping: Dict[str, Any]) -> bool:
        """
        批量设置

        :param mapping: 键值字典
        :return: 是否成功
        """
        return self.client.mset(mapping)

    def mget(self, keys: List[str]) -> List[Any]:
        """
        批量获取

        :param keys: 键列表
        :return: 值列表
        """
        return self.client.mget(keys)

    def incr(self, key: str, amount: int = 1) -> int:
        """
        自增

        :param key: 键
        :param amount: 增量
        :return: 自增后的值
        """
        return self.client.incr(key, amount=amount)

    def decr(self, key: str, amount: int = 1) -> int:
        """
        自减

        :param key: 键
        :param amount: 减量
        :return: 自减后的值
        """
        return self.client.decr(key, amount=amount)

    # ============ List 操作 ============

    def lpush(self, name: str, *values: Any) -> int:
        """
        从列表左侧添加元素

        :param name: 列表名
        :param values: 值列表
        :return: 列表长度
        """
        return self.client.lpush(name, *values)

    def rpush(self, name: str, *values: Any) -> int:
        """
        从列表右侧添加元素

        :param name: 列表名
        :param values: 值列表
        :return: 列表长度
        """
        return self.client.rpush(name, *values)

    def lrange(self, name: str, start: int, end: int) -> List[Any]:
        """
        获取列表元素

        :param name: 列表名
        :param start: 起始索引
        :param end: 结束索引
        :return: 元素列表
        """
        return self.client.lrange(name, start, end)

    def lpop(self, name: str) -> Optional[Any]:
        """
        从列表左侧移出元素

        :param name: 列表名
        :return: 移出的元素
        """
        return self.client.lpop(name)

    def rpop(self, name: str) -> Optional[Any]:
        """
        从列表右侧移出元素

        :param name: 列表名
        :return: 移出的元素
        """
        return self.client.rpop(name)

    def brpop(self, name: str, timeout: int = 0) -> Optional[Tuple[str, Any]]:
        """
        从列表右侧移出元素（阻塞）

        :param name: 列表名
        :param timeout: 超时时间（秒）
        :return: 移出的元素
        """
        return self.client.brpop(name, timeout=timeout)

    def llen(self, name: str) -> int:
        """
        获取列表长度

        :param name: 列表名
        :return: 列表长度
        """
        return self.client.llen(name)

    # ============ Set 操作 ============

    def sadd(self, name: str, *values: Any) -> int:
        """
        向集合添加元素

        :param name: 集合名
        :param values: 值列表
        :return: 添加的元素数量
        """
        return self.client.sadd(name, *values)

    def smembers(self, name: str) -> set[Any]:
        """
        获取集合所有元素

        :param name: 集合名
        :return: 元素集合
        """
        return self.client.smembers(name)

    def sinter(self, keys: List[str]) -> set[Any]:
        """
        获取多个集合的交集

        :param keys: 集合名列表
        :return: 交集
        """
        return self.client.sinter(keys)

    def sunion(self, keys: List[str]) -> set[Any]:
        """
        获取多个集合的并集

        :param keys: 集合名列表
        :return: 并集
        """
        return self.client.sunion(keys)

    def srem(self, name: str, *values: Any) -> int:
        """
        从集合中移除元素

        :param name: 集合名
        :param values: 值列表
        :return: 移除的元素数量
        """
        return self.client.srem(name, *values)

    def scard(self, name: str) -> int:
        """
        获取集合元素个数

        :param name: 集合名
        :return: 元素个数
        """
        return self.client.scard(name)

    # ============ Sorted Set 操作 ============

    def zadd(self, name: str, mapping: Dict[Any, float]) -> int:
        """
        向有序集合添加元素

        :param name: 有序集合名
        :param mapping: 元素-分数字典
        :return: 添加的元素数量
        """
        return self.client.zadd(name, mapping)

    def zrange(
        self,
        name: str,
        start: int,
        end: int,
        withscores: bool = False,
        desc: bool = False,
    ) -> Union[List[Any], List[Tuple[Any, float]]]:
        """
        获取有序集合元素

        :param name: 有序集合名
        :param start: 起始索引
        :param end: 结束索引
        :param withscores: 是否包含分数
        :param desc: 是否降序
        :return: 元素列表
        """
        return self.client.zrange(name, start, end, withscores=withscores, desc=desc)

    def zrevrange(
        self, name: str, start: int, end: int, withscores: bool = False
    ) -> Union[List[Any], List[Tuple[Any, float]]]:
        """
        获取有序集合元素（降序）

        :param name: 有序集合名
        :param start: 起始索引
        :param end: 结束索引
        :param withscores: 是否包含分数
        :return: 元素列表
        """
        return self.client.zrevrange(name, start, end, withscores=withscores)

    def zrangebyscore(
        self,
        name: str,
        min_score: float,
        max_score: float,
        withscores: bool = False,
    ) -> Union[List[Any], List[Tuple[Any, float]]]:
        """
        按分数范围获取有序集合元素

        :param name: 有序集合名
        :param min_score: 最小分数
        :param max_score: 最大分数
        :param withscores: 是否包含分数
        :return: 元素列表
        """
        return self.client.zrangebyscore(name, min_score, max_score, withscores=withscores)

    def zremrangebyrank(self, name: str, min_rank: int, max_rank: int) -> int:
        """
        按排名范围移除有序集合元素

        :param name: 有序集合名
        :param min_rank: 最小排名
        :param max_rank: 最大排名
        :return: 移除的元素数量
        """
        return self.client.zremrangebyrank(name, min_rank, max_rank)

    def zcard(self, name: str) -> int:
        """
        获取有序集合元素个数

        :param name: 有序集合名
        :return: 元素个数
        """
        return self.client.zcard(name)

    # ============ Hash 操作 ============

    def hset(self, name: str, key: str, value: Any) -> int:
        """
        设置哈希字段

        :param name: 哈希名
        :param key: 字段名
        :param value: 值
        :return: 是否新增
        """
        return self.client.hset(name, key, value)

    def hget(self, name: str, key: str) -> Optional[Any]:
        """
        获取哈希字段值

        :param name: 哈希名
        :param key: 字段名
        :return: 值
        """
        return self.client.hget(name, key)

    def hgetall(self, name: str) -> Dict[str, Any]:
        """
        获取哈希所有字段

        :param name: 哈希名
        :return: 字段字典
        """
        return self.client.hgetall(name)

    def hmset(self, name: str, mapping: Dict[str, Any]) -> bool:
        """
        批量设置哈希字段

        :param name: 哈希名
        :param mapping: 字段字典
        :return: 是否成功
        """
        return self.client.hmset(name, mapping)

    def hdel(self, name: str, *keys: str) -> int:
        """
        删除哈希字段

        :param name: 哈希名
        :param keys: 字段名列表
        :return: 删除的字段数量
        """
        return self.client.hdel(name, *keys)

    # ============ 通用操作 ============

    def delete(self, *keys: str) -> int:
        """
        删除键

        :param keys: 键列表
        :return: 删除的键数量
        """
        return self.client.delete(*keys)

    def exists(self, *keys: str) -> int:
        """
        判断键是否存在

        :param keys: 键列表
        :return: 存在的键数量
        """
        return self.client.exists(*keys)

    def expire(self, key: str, time: int) -> bool:
        """
        设置键的过期时间

        :param key: 键
        :param time: 过期时间（秒）
        :return: 是否成功
        """
        return self.client.expire(key, time)

    def ttl(self, key: str) -> int:
        """
        获取键的剩余过期时间

        :param key: 键
        :return: 剩余时间（秒），-1 表示永久，-2 表示不存在
        """
        return self.client.ttl(key)

    def scan_iter(self, match: str = "*", count: int = 10) -> Any:
        """
        扫描键

        :param match: 匹配模式
        :param count: 每次扫描数量
        :return: 键迭代器
        """
        return self.client.scan_iter(match=match, count=count)

    def scan_del(self, match: str = "djes:", size: int = 10000) -> int:
        """
        扫描并删除匹配的键

        :param match: 匹配模式
        :param size: 批量删除大小
        :return: 删除的总数
        """
        pipe = self.client.pipeline()
        deleted_count = 0
        for i, key in enumerate(self.client.scan_iter(match=match + "*", count=1000), 1):
            pipe.delete(key)
            if i % size == 0:
                pipe.execute()
                print(f"已删除 {i} 个键")
            deleted_count += 1
        pipe.execute()
        return deleted_count

    # ============ Lua 脚本 ============

    def register_script(self, script: str) -> Any:
        """
        注册 Lua 脚本

        :param script: Lua 脚本内容
        :return: 脚本对象
        """
        return self.client.register_script(script)

    # ============ 上下文管理器 ============

    def __enter__(self) -> "RedisClient":
        """支持上下文管理器"""
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """退出上下文"""
        pass


if __name__ == "__main__":
    # 单机模式示例
    rc = RedisClient(host="127.0.0.1", port=6379, db=0)

    # String 操作
    rc.set("test", "test")
    rc.set("test_num", 5)
    print("get test:", rc.get("test"))
    print("mset:", rc.mset({"test1": "test1", "test2": "test2"}))
    print("mget:", rc.mget(["test1", "test2"]))
    print("incr:", rc.incr("test_num"))
    print("decr:", rc.decr("test_num"))

    # List 操作
    rc.delete("test_list")
    rc.lpush("test_list", 1, 2, 3)
    print("lrange:", rc.lrange("test_list", 0, 10))
    print("brpop:", rc.brpop("test_list"))
    print("llen:", rc.llen("test_list"))

    # Set 操作
    rc.sadd("test_set", *[1, 2, 3, 4])
    print("smembers:", rc.smembers("test_set"))
    rc.sadd("test_set2", *[3, 4, 5, 6])
    print("sinter:", rc.sinter(["test_set", "test_set2"]))
    print("sunion:", rc.sunion(["test_set", "test_set2"]))
    print("scard:", rc.scard("test_set"))

    # Sorted Set 操作
    rc.zadd("test_zset", {"a": 1, "b": 10})
    rc.zadd("test_zset1", {"b": 2, "c": 20})
    print("zrange:", rc.zrange("test_zset", 0, -1, withscores=True))
    print("zrevrange:", rc.zrevrange("test_zset", 0, -1, withscores=True))
    print("zrangebyscore:", rc.zrangebyscore("test_zset", 0, 5, withscores=True))

    # Hash 操作
    rc.hmset("test_hash", {"a": 1, "b": 10})
    print("hgetall:", rc.hgetall("test_hash"))

    # 集群模式示例（需要配置真实的集群节点）
    # startup_nodes = [{"host": "10.10.20.97", "port": 7000}, {"host": "10.10.20.97", "port": 7001}]
    # rc_cluster = RedisClient(startup_nodes=startup_nodes)

    # 使用上下文管理器
    with RedisClient(host="127.0.0.1", port=6379) as redis_ctx:
        redis_ctx.set("ctx_test", "value")
        print("ctx_test:", redis_ctx.get("ctx_test"))
