import uuid

from redis.exceptions import (
    BusyLoadingError,
    ConnectionError,
    TimeoutError,
    TryAgainError,
)
from redis.sentinel import MasterNotFoundError

from pypeline.utils.redis_client import ResilientRedisClient


TRANSIENT_REDIS_ERRORS = (
    ConnectionError,
    TimeoutError,
    MasterNotFoundError,
    BusyLoadingError,
    TryAgainError,
)


class LockingParallelBarrier:
    def __init__(self, redis_url, task_key="task_counter", lock_key="task_lock"):
        self.redis_client = ResilientRedisClient(redis_url, decode_responses=True)
        self.redis = self.redis_client.redis

        self.task_key = task_key
        self.lock_key = lock_key
        self.lock_token = None

        self._release_lock_script = self.redis.register_script(
            """
            if redis.call("get", KEYS[1]) == ARGV[1] then
                return redis.call("del", KEYS[1])
            else
                return 0
            end
            """
        )

    def _sleep_with_exponential_backoff(
        self, attempt: int, retry_interval: float
    ) -> None:
        self.redis_client._sleep_with_exponential_backoff(attempt, retry_interval)

    def _call_with_failover_retry(
        self,
        fn,
        *,
        timeout: float = 180.0,
        retry_interval: float = 0.25,
        op_name: str = "redis operation",
        retry_on_result=None,
    ):
        deadline = __import__("time").monotonic() + timeout
        attempt = 0

        while True:
            try:
                result = fn()

                if retry_on_result is not None and retry_on_result(result):
                    if __import__("time").monotonic() >= deadline:
                        raise TimeoutError(
                            f"{op_name} timed out after {timeout}s without success"
                        )
                    self._sleep_with_exponential_backoff(attempt, retry_interval)
                    attempt += 1
                    continue

                return result

            except TRANSIENT_REDIS_ERRORS as exc:
                if __import__("time").monotonic() >= deadline:
                    raise TimeoutError(
                        f"{op_name} failed after {timeout}s; last error: {exc}"
                    ) from exc

                self._sleep_with_exponential_backoff(attempt, retry_interval)
                attempt += 1

    def advance_once(self, advance_key: str, ttl: int) -> bool:
        """
        Returns True exactly once per advance_key (within ttl).
        Subsequent calls return False.
        """
        return bool(
            self._call_with_failover_retry(
                lambda: self.redis.set(advance_key, "1", nx=True, ex=ttl),
                op_name=f"advance_once({advance_key})",
            )
        )

    def acquire_lock(self, timeout=180, retry_interval=0.25):
        """
        Acquire a lock using Redis and store an ownership token.

        Preserves the existing public signature:
        - timeout: total wait budget AND lock TTL
        - retry_interval: base interval for exponential backoff

        Raises TimeoutError if the lock cannot be acquired within timeout seconds.
        """
        token = str(uuid.uuid4())

        result = self._call_with_failover_retry(
            lambda: self.redis.set(self.lock_key, token, nx=True, ex=timeout),
            timeout=timeout,
            retry_interval=retry_interval,
            op_name=f"acquire_lock({self.lock_key})",
            retry_on_result=lambda result: not bool(result),
        )

        if result:
            self.lock_token = token
            return True

        raise TimeoutError(
            f"Could not acquire lock '{self.lock_key}' within {timeout}s"
        )

    def release_lock(self):
        """Release the lock only if this instance still owns it."""
        if not self.lock_token:
            return False

        token = self.lock_token
        try:
            released = self._call_with_failover_retry(
                lambda: self._release_lock_script(
                    keys=[self.lock_key],
                    args=[token],
                ),
                op_name=f"release_lock({self.lock_key})",
            )
            return bool(released)
        finally:
            self.lock_token = None

    def set_task_count(self, count):
        """Initialize the task counter in Redis."""
        self._call_with_failover_retry(
            lambda: self.redis.set(self.task_key, count),
            op_name=f"set_task_count({self.task_key})",
        )

    def decrement_task_count(self):
        """Decrement the task counter in Redis."""
        return int(
            self._call_with_failover_retry(
                lambda: self.redis.decr(self.task_key),
                op_name=f"decrement_task_count({self.task_key})",
            )
        )

    def task_exists(self):
        return bool(
            self._call_with_failover_retry(
                lambda: self.redis.exists(self.task_key),
                op_name=f"task_exists({self.task_key})",
            )
        )

    def get_task_count(self):
        """Get the current value of the task counter."""
        value = self._call_with_failover_retry(
            lambda: self.redis.get(self.task_key),
            op_name=f"get_task_count({self.task_key})",
        )
        return int(value or 0)
