# -*- coding: utf-8 -*-
"""
@File    : distributed_lock.py
@Author  : Jay
@Date    : 2026/3/18 17:47
@Desc    : Redis connection pool and distributed lock implementation
"""
import asyncio
import time
import uuid

_ACQUIRE_LUA = """
local ttl_ms = tonumber(ARGV[#ARGV])
for i=1,#KEYS do
    if redis.call('exists', KEYS[i]) == 1 then
        return 0
    end
    redis.call('set', KEYS[i], ARGV[i], 'PX', ttl_ms)
end
return 1
"""
_RELEASE_LUA = """
for i=1,#KEYS do
  if redis.call('get', KEYS[i]) == ARGV[i] then
    redis.call('del', KEYS[i])
  end
end
return 1
"""


# Distributed lock class, use `aioredlock` for cluster-wide locks
class AsyncRedisAtomicMultiLock:
    """Implements a Redis-based atomic multi-key lock with optional blocking behavior."""

    def __init__(
            self,
            *keys: str,
            redis_xsync,
            redis_db,
            ttl_ms: int = 5000,
            blocking: bool = False,
            blocking_timeout: int = 5,
            sleep: float = 0.3
    ):
        """Initialize the atomic multi-key lock.

        Args:
            *keys (str): Redis keys to lock, can pass multiple keys.
            redis_xsync: Redis client factory supporting sync/async context.
            redis_db: Redis database number.
            ttl_ms (int): Lock timeout in milliseconds.
            blocking (bool): If True, wait until lock is acquired when already locked.
            blocking_timeout (int): Maximum blocking wait time in seconds.
            sleep (float): Sleep interval between retries when blocking.
        """
        self.redis_xsync = redis_xsync
        self.redis_db = redis_db
        self.keys = list(keys)
        self.ttl_ms = int(ttl_ms)
        self.container = list(map(lambda _: uuid.uuid4().hex, keys))  # Unique lock values
        self.acquired = False

        self.blocking = blocking
        self.blocking_timeout = blocking_timeout
        self.sleep = sleep

    # <editor-fold desc="Distributed lock: async implementation">
    async def _async_with_client(self, script: str, keys, *args):
        """Helper to run Lua scripts asynchronously against Redis.

        Args:
            script (str): Lua script to execute.
            keys (list): List of Redis keys.
            *args: Arguments for the Lua script.

        Returns:
            bool: True if Lua script returned success (1), False otherwise.
        """
        async with self.redis_xsync(db=self.redis_db) as redis:
            return (await redis.eval(script, len(keys), *keys, *args)) == 1

    async def async_acquire(self) -> bool:
        """Asynchronously acquire the lock.

        Returns:
            bool: True if lock acquired successfully, False otherwise.
        """
        argv = [*self.container, str(self.ttl_ms)]

        async def try_once():
            return await self._async_with_client(_ACQUIRE_LUA, self.keys, *argv)

        if not self.blocking:
            self.acquired = await try_once()
            return self.acquired

        end_time = asyncio.get_event_loop().time() + self.blocking_timeout
        while True:
            if await try_once():
                self.acquired = True
                return True
            if asyncio.get_event_loop().time() > end_time:
                return False
            await asyncio.sleep(self.sleep)

    async def async_release(self) -> None:
        """Asynchronously release the lock, only if held by this instance."""
        try:
            await self._async_with_client(_RELEASE_LUA, self.keys, *self.container)
        finally:
            self.acquired = False

    async def __aenter__(self):
        """Async context manager entry.

        Raises:
            TimeoutError: If lock cannot be acquired within timeout.

        Returns:
            AsyncRedisAtomicMultiLock: self
        """
        ok = await self.async_acquire()
        if not ok:
            raise TimeoutError("Failed to acquire lock within timeout")
        return self

    async def __aexit__(self, exc_type, exc, tb):
        """Async context manager exit."""
        await self.async_release()

    # </editor-fold>

    # <editor-fold desc="Distributed lock: sync implementation">
    def _sync_with_client(self, script: str, keys, *args) -> bool:
        """Helper to run Lua scripts synchronously against Redis.

        Args:
            script (str): Lua script to execute.
            keys (list): List of Redis keys.
            *args: Arguments for the Lua script.

        Returns:
            bool: True if Lua script returned success (1), False otherwise.
        """
        with self.redis_xsync(db=self.redis_db) as redis:
            return int(redis.eval(script, len(keys), *keys, *args)) == 1

    def sync_acquire(self) -> bool:
        """Synchronously acquire the lock.

        Returns:
            bool: True if lock acquired successfully, False otherwise.
        """
        argv = [*self.container, str(self.ttl_ms)]

        def try_once():
            return self._sync_with_client(_ACQUIRE_LUA, self.keys, *argv)

        if not self.blocking:
            self.acquired = try_once()
            return self.acquired

        end_time = time.time() + self.blocking_timeout
        while True:
            if try_once():
                self.acquired = True
                return True
            if time.time() > end_time:
                return False
            time.sleep(self.sleep)

    def sync_release(self) -> None:
        """Synchronously release the lock."""
        try:
            self._sync_with_client(_RELEASE_LUA, self.keys, *self.container)
        finally:
            self.acquired = False

    def __enter__(self):
        """Sync context manager entry.

        Raises:
            TimeoutError: If lock cannot be acquired within timeout.

        Returns:
            AsyncRedisAtomicMultiLock: self
        """
        ok = self.sync_acquire()
        if not ok:
            raise TimeoutError(
                f"Failed to acquire lock on keys {self.keys} in {self.blocking_timeout}s"
            )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Sync context manager exit."""
        self.sync_release()
    # </editor-fold>


__all__ = ["AsyncRedisAtomicMultiLock"]
