import time
import uuid
import redis
from redis.backoff import ExponentialBackoff
from redis.retry import Retry
from redis.exceptions import ConnectionError, TimeoutError
from redis.sentinel import Sentinel
from urllib.parse import urlparse

from pypeline.constants import (
    REDIS_SENTINEL_MASTER_NAME,
    DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
    DEFAULT_REDIS_SOCKET_TIMEOUT,
    DEFAULT_REDIS_RETRY_ON_TIMEOUT,
    DEFAULT_REDIS_SOCKET_KEEPALIVE,
    DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
)


class LockingParallelBarrier:
    def __init__(self, redis_url, task_key="task_counter", lock_key="task_lock"):
        retry = Retry(ExponentialBackoff(), 5)

        common_kwargs = {
            "socket_connect_timeout": DEFAULT_REDIS_SOCKET_CONNECT_TIMEOUT,
            "socket_timeout": DEFAULT_REDIS_SOCKET_TIMEOUT,
            "retry_on_timeout": DEFAULT_REDIS_RETRY_ON_TIMEOUT,
            "socket_keepalive": DEFAULT_REDIS_SOCKET_KEEPALIVE,
            "health_check_interval": DEFAULT_REDIS_HEALTH_CHECK_INTERVAL,
            "retry": retry,
            "retry_on_error": [ConnectionError, TimeoutError],
        }

        if REDIS_SENTINEL_MASTER_NAME is not None:
            parsed_redis_url = urlparse(redis_url)
            redis_sentinel = Sentinel(
                sentinels=[(parsed_redis_url.hostname, parsed_redis_url.port)],
                sentinel_kwargs=common_kwargs,
            )
            self.redis = redis_sentinel.master_for(
                REDIS_SENTINEL_MASTER_NAME,
                db=int(parsed_redis_url.path[1]) if parsed_redis_url.path else 0,
                password=parsed_redis_url.password,
                decode_responses=True,
                **common_kwargs,
            )
        else:
            self.redis = redis.StrictRedis.from_url(
                redis_url,
                decode_responses=True,
                **common_kwargs,
            )

        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 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.redis.set(advance_key, "1", nx=True, ex=ttl))

    def acquire_lock(self, timeout=5, retry_interval=0.1):
        """Acquire a lock using Redis and store an ownership token."""
        token = str(uuid.uuid4())

        while True:
            if self.redis.set(self.lock_key, token, nx=True, ex=timeout):
                self.lock_token = token
                return True
            time.sleep(retry_interval)

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

        try:
            released = self._release_lock_script(
                keys=[self.lock_key],
                args=[self.lock_token],
            )
            return bool(released)
        finally:
            self.lock_token = None

    def set_task_count(self, count):
        """Initialize the task counter in Redis."""
        self.redis.set(self.task_key, count)

    def decrement_task_count(self):
        """Decrement the task counter in Redis."""
        return self.redis.decr(self.task_key)

    def task_exists(self):
        return self.redis.exists(self.task_key)

    def get_task_count(self):
        """Get the current value of the task counter."""
        return int(self.redis.get(self.task_key) or 0)
