from functools import wraps

from redis_lock import Lock as RedisLock
from redis import Redis

from common.utils import get_logger
from common.utils.inspect import copy_function_args
from apps.jumpserver.const import CONFIG

logger = get_logger(__file__)


class AcquireFailed(RuntimeError):
    pass


class DistributedLock(RedisLock):
    def __init__(self, name, blocking=True, expire=60*2, auto_renewal=True):
        """
        使用 redis 构造的分布式锁

        :param name:
            锁的名字，要全局唯一
        :param blocking:
            该参数只在锁作为装饰器或者 `with` 时有效。
        :param expire:
            锁的过期时间，注意不一定是锁到这个时间就释放了，分两种情况
            当 `auto_renewal=False` 时，锁会释放
            当 `auto_renewal=True` 时，如果过期之前程序还没释放锁，我们会延长锁的存活时间。
            这里的作用是防止程序意外终止没有释放锁，导致死锁。
        """
        self.kwargs_copy = copy_function_args(self.__init__, locals())
        redis = Redis(host=CONFIG.REDIS_HOST, port=CONFIG.REDIS_PORT, password=CONFIG.REDIS_PASSWORD)
        super().__init__(redis_client=redis, name=name, expire=expire, auto_renewal=auto_renewal)
        self._blocking = blocking

    def __enter__(self):
        acquired = self.acquire(blocking=self._blocking)
        if self._blocking and not acquired:
            raise EnvironmentError("Lock wasn't acquired, but blocking=True")
        if not acquired:
            raise AcquireFailed
        return self

    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        self.release()

    def __call__(self, func):
        @wraps(func)
        def inner(*args, **kwds):
            # 要创建一个新的锁对象
            with self.__class__(**self.kwargs_copy):
                return func(*args, **kwds)

        return inner