# !/usr/bin/env python
# -*-coding:utf-8 -*-

"""
基于 Redis 的分布式锁实现

用于确保定时任务在多实例部署时只在一个实例上执行。

特性:
- 基于 Redis SET NX EX 实现原子性锁
- 支持锁超时自动释放
- 支持锁延期（防止长任务锁过期）
- 使用 Lua 脚本确保删除安全性
- 支持装饰器和上下文管理器两种使用方式
"""

import functools
import uuid
from typing import Optional, Callable, Any, TYPE_CHECKING

from loguru import logger

if TYPE_CHECKING:
    from redis.asyncio import Redis


class RedisDistributedLock:
    """
    Redis 分布式锁

    使用 Redis 的 SET NX EX 命令实现分布式锁，确保多实例环境下的互斥访问。
    """

    def __init__(
        self,
        redis_client: "Redis",
        key: str,
        timeout: int = 300,
        identifier: Optional[str] = None
    ):
        """
        初始化分布式锁

        Args:
            redis_client: Redis 客户端实例
            key: 锁的键名（自动添加 'lock:' 前缀）
            timeout: 锁超时时间（秒），默认 300 秒
            identifier: 锁标识符（用于确保只能释放自己的锁），默认随机生成
        """
        self.redis = redis_client
        self.key = f"lock:{key}"
        self.timeout = timeout
        self.identifier = identifier or str(uuid.uuid4())
        self._acquired = False

    async def acquire(self) -> bool:
        """
        获取锁

        使用 SET NX EX 命令原子性地获取锁。

        Returns:
            是否成功获取锁
        """
        try:
            # SET key value NX EX timeout
            # NX: 只在键不存在时设置
            # EX: 设置过期时间（秒）
            result = await self.redis.set(
                self.key,
                self.identifier,
                ex=self.timeout,
                nx=True
            )

            if result:
                self._acquired = True
                logger.debug(f"🔒 获取锁成功: {self.key} (id: {self.identifier[:8]})")
                return True
            else:
                logger.debug(f"🔒 锁被占用: {self.key}")
                return False

        except Exception as e:
            logger.error(f"❌ 获取锁失败: {self.key}, 错误: {e}")
            return False

    async def release(self) -> bool:
        """
        释放锁

        使用 Lua 脚本确保只删除自己持有的锁，防止误删其他实例的锁。

        Returns:
            是否成功释放锁
        """
        if not self._acquired:
            logger.debug(f"⚠️ 锁未获取，无需释放: {self.key}")
            return False

        try:
            # Lua 脚本：只有当锁的值等于自己的标识符时才删除
            lua_script = """
            if redis.call("get", KEYS[1]) == ARGV[1] then
                return redis.call("del", KEYS[1])
            else
                return 0
            end
            """

            result = await self.redis.eval(lua_script, 1, self.key, self.identifier)

            if result:
                self._acquired = False
                logger.debug(f"🔓 释放锁成功: {self.key} (id: {self.identifier[:8]})")
                return True
            else:
                logger.warning(
                    f"⚠️ 释放锁失败（可能已过期或被其他进程持有）: {self.key}"
                )
                return False

        except Exception as e:
            logger.error(f"❌ 释放锁异常: {self.key}, 错误: {e}")
            return False

    async def extend(self, additional_time: Optional[int] = None) -> bool:
        """
        延长锁的过期时间

        用于长时间运行的任务，防止锁在任务完成前过期。

        Args:
            additional_time: 额外延长的时间（秒），默认使用初始 timeout

        Returns:
            是否成功延长
        """
        if not self._acquired:
            logger.debug(f"⚠️ 锁未获取，无法延期: {self.key}")
            return False

        try:
            extend_time = additional_time or self.timeout

            # Lua 脚本：只有当锁的值等于自己的标识符时才延期
            lua_script = """
            if redis.call("get", KEYS[1]) == ARGV[1] then
                return redis.call("expire", KEYS[1], ARGV[2])
            else
                return 0
            end
            """

            result = await self.redis.eval(
                lua_script,
                1,
                self.key,
                self.identifier,
                extend_time
            )

            if result:
                logger.debug(f"⏰ 延长锁成功: {self.key}, 延长 {extend_time} 秒")
                return True
            else:
                logger.warning(f"⚠️ 延长锁失败: {self.key}")
                return False

        except Exception as e:
            logger.error(f"❌ 延长锁异常: {self.key}, 错误: {e}")
            return False

    async def __aenter__(self):
        """异步上下文管理器：进入"""
        await self.acquire()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器：退出"""
        await self.release()


def distributed_task(
    lock_key: Optional[str] = None,
    timeout: int = 300,
    skip_locked: bool = True
):
    """
    分布式任务装饰器

    确保定时任务在多实例部署时只在一个实例上执行。

    Args:
        lock_key: 锁的键名（默认使用函数名）
        timeout: 锁超时时间（秒）
        skip_locked: 获取锁失败时是否跳过执行（True）还是等待（False）

    Example:
        @distributed_task(timeout=300)
        async def my_scheduled_task():
            # 任务逻辑
            # 在多实例部署时，同一时刻只有一个实例会执行
            pass

    Usage with scheduler:
        @cron(hour=0, minute=0)
        @distributed_task(lock_key="daily_cleanup", timeout=600)
        async def daily_cleanup():
            # 每天 00:00 执行，多实例环境下只有一个实例执行
            pass
    """
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        async def wrapper(*args, **kwargs) -> Any:
            # 生成锁键名
            key = lock_key or f"task:{func.__name__}"

            # 获取 Redis 客户端
            try:
                # 尝试从 app.state 获取 Redis
                from infoman.service.app import application
                if not hasattr(application.state, 'redis'):
                    logger.warning(
                        f"⚠️ Redis 未初始化，任务 {func.__name__} 将直接执行（无分布式锁保护）"
                    )
                    return await func(*args, **kwargs)

                redis_client = application.state.redis

            except Exception as e:
                logger.warning(
                    f"⚠️ 获取 Redis 客户端失败: {e}，任务 {func.__name__} 将直接执行"
                )
                return await func(*args, **kwargs)

            # 创建分布式锁
            lock = RedisDistributedLock(redis_client, key, timeout)

            # 尝试获取锁
            acquired = await lock.acquire()

            if not acquired:
                if skip_locked:
                    logger.info(
                        f"⏭️ 任务 {func.__name__} 跳过执行（锁被其他实例占用）"
                    )
                    return None
                else:
                    logger.info(
                        f"⏳ 任务 {func.__name__} 等待锁释放..."
                    )
                    # TODO: 实现等待逻辑（可选）
                    return None

            # 执行任务
            try:
                logger.info(f"▶️ 任务 {func.__name__} 开始执行 (lock_id: {lock.identifier[:8]})")
                result = await func(*args, **kwargs)
                logger.success(f"✅ 任务 {func.__name__} 执行完成")
                return result

            except Exception as e:
                logger.error(f"❌ 任务 {func.__name__} 执行失败: {e}", exc_info=True)
                raise

            finally:
                # 释放锁
                await lock.release()

        return wrapper
    return decorator


async def check_redis_connection(redis_client: "Redis") -> bool:
    """
    检查 Redis 连接是否正常

    Args:
        redis_client: Redis 客户端

    Returns:
        连接是否正常
    """
    try:
        await redis_client.ping()
        logger.info("✅ Redis 连接正常")
        return True
    except Exception as e:
        logger.error(f"❌ Redis 连接异常: {e}")
        return False
