"""Distributed lock protocol and implementations for fair task queue."""

from contextlib import AbstractAsyncContextManager
from datetime import timedelta
from typing import Protocol


class DistributedLock(Protocol):
    """Protocol for distributed locks.

    Distributed locks are used to coordinate operations across multiple
    scheduler instances in a distributed deployment.
    """

    async def acquire(
        self,
        lock_key: str,
        timeout_duration: timedelta,
        lock_ttl: timedelta | None = None,
    ) -> bool:
        """Acquire a distributed lock.

        Args:
            lock_key: Unique key for the lock
            timeout_duration: Maximum time to wait for lock acquisition
            lock_ttl: Optional TTL for the lock (auto-release after this time)

        Returns:
            True if lock was acquired, False if timeout was reached

        Raises:
            FairQueueStorageError: If lock operation fails

        """
        ...

    async def release(self, lock_key: str) -> None:
        """Release a distributed lock.

        Args:
            lock_key: Key of the lock to release

        Raises:
            FairQueueStorageError: If lock release fails

        """
        ...

    async def is_locked(self, lock_key: str) -> bool:
        """Check if a lock is currently held.

        Args:
            lock_key: Key of the lock to check

        Returns:
            True if lock is held, False otherwise

        Raises:
            FairQueueStorageError: If check operation fails

        """
        ...


class DistributedLockContext(AbstractAsyncContextManager[str]):
    """Context manager for distributed locks.

    Example:
        ```python
        async with DistributedLockContext(lock, "recovery_lock", timeout_duration=timedelta(seconds=5)):
            # Lock is acquired, perform operations
            await recover_stuck_tasks()
        # Lock is automatically released
        ```

    """

    def __init__(
        self,
        lock: DistributedLock,
        lock_key: str,
        timeout_duration: timedelta,
        lock_ttl: timedelta | None = None,
    ) -> None:
        """Initialize distributed lock context.

        Args:
            lock: Distributed lock implementation
            lock_key: Unique key for the lock
            timeout_duration: Maximum time to wait for lock acquisition
            lock_ttl: Optional TTL for the lock

        """
        self._lock = lock
        self._lock_key = lock_key
        self._timeout_duration = timeout_duration
        self._lock_ttl = lock_ttl
        self._acquired = False

    async def __aenter__(self) -> str:
        """Acquire lock and return lock key."""
        self._acquired = await self._lock.acquire(
            self._lock_key,
            self._timeout_duration,
            self._lock_ttl,
        )
        if not self._acquired:
            lock_key = self._lock_key
            timeout_duration = self._timeout_duration
            msg = f"Failed to acquire lock {lock_key} within {timeout_duration}"
            raise TimeoutError(msg)
        return self._lock_key

    async def __aexit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
    ) -> None:
        """Release lock if acquired."""
        if self._acquired:
            await self._lock.release(self._lock_key)
