import logging
import json
import uuid
import hashlib
import redis
from celery import Task

from rhodecode import CONFIG
from rhodecode.lib.celerylib import RequestContextTask

log = logging.getLogger(__name__)


def stable_payload_key(task_name: str, args: tuple, kwargs: dict) -> str:
    """Deterministic key from (name, args, kwargs) ignoring any private kwargs."""
    clean_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
    payload = json.dumps([task_name, args, clean_kwargs], sort_keys=True, separators=(",", ":"))
    digest = hashlib.sha256(payload.encode()).hexdigest()
    return f"celery:unique:{task_name}:{digest}"


class UniqueByPayloadTask(Task):
    """Base task that deduplicates by (task name + args + kwargs) at enqueue time."""

    abstract = True

    _redis_url = None
    _lock_ttl_seconds = None
    _redis_connection = None
    _deduplication_enabled = None
    _redis_available = None

    @classmethod
    def redis_url(cls):
        if cls._redis_url is None:
            cls._redis_url = CONFIG["celery.deduplicate_lock_store"]
        return cls._redis_url

    @classmethod
    def _redis(cls):
        if cls._redis_connection is None:
            log.debug("Connecting to Redis at %s", cls.redis_url())
            cls._redis_connection = redis.from_url(cls.redis_url())
        return cls._redis_connection

    @property
    def deduplication_enabled(self):
        if self._deduplication_enabled is None:
            self._deduplication_enabled = CONFIG["celery.enable_deduplication"]
        return self._deduplication_enabled

    @property
    def lock_ttl_seconds(self):
        if self._lock_ttl_seconds is None:
            self._lock_ttl_seconds = CONFIG["celery.lock_ttl_seconds"]
        return self._lock_ttl_seconds

    def _unique_key(self, args, kwargs) -> str:
        return stable_payload_key(self.name, args, kwargs)

    def apply_async(
        self, args=None, kwargs=None, task_id=None, producer=None, link=None, link_error=None, shadow=None, **options
    ):
        """
        If another identical (name+payload) task is queued/running,
        return its AsyncResult instead of enqueueing a new message.
        """

        self._validate_config_and_redis_available()

        if not self.deduplication_enabled:
            log.debug("Deduplication disabled; skipping unique key check")
            return super().apply_async(args, kwargs, task_id, producer, link, link_error, shadow, **options)

        args = args or ()
        kwargs = kwargs or {}

        key = self._unique_key(args, kwargs)
        r = self._redis()

        # Generate a task_id up front so we can store/return it consistently.
        task_id = uuid.uuid4().hex

        # Atomically claim the key if absent
        claimed = r.set(key, task_id, nx=True, ex=self.lock_ttl_seconds)
        log.debug("Task %s unique key claim: %s", key, claimed)
        if not claimed:
            # Someone else already queued (or is running) the same task.
            from celery.result import AsyncResult

            existing = r.get(key)
            log.debug("Task %s already queued or running; returning existing result", key)
            return AsyncResult((existing or b"").decode() or None, app=self._get_app())

        options["headers"] = options.get("headers", {})
        options["headers"].update({"unique_key": key})

        return super().apply_async(args, kwargs, task_id, producer, link, link_error, shadow, **options)

    def after_return(self, status, retval, task_id, args, kwargs, einfo):
        if not self.deduplication_enabled:
            log.debug("Deduplication disabled; skipping unique key check")
            return super().after_return(status, retval, task_id, args, kwargs, einfo)

        key = None
        try:
            key = getattr(self.request, "headers", {}).get("unique_key")
            log.debug("Deleting lock %s", key)
            if key:
                self._redis().delete(key)
        except Exception as e:
            # Don't crash the worker if Redis is momentarily unavailable.
            if key:
                log.warning("Failed to delete lock %s: %s", key, e)

        return super().after_return(status, retval, task_id, args, kwargs, einfo)

    def _validate_config_and_redis_available(self):
        if (
            CONFIG.get("celery.enable_deduplication") is None
            or CONFIG.get("celery.deduplicate_lock_store") is None
            or CONFIG.get("celery.lock_ttl_seconds") is None
        ):
            log.warning("Missing config; unique task deduplication disabled")
            self._deduplication_enabled = False
            self._redis_available = False

        if self._redis_available is None:
            try:
                self._redis()
                self._redis_available = True
            except Exception as e:
                log.warning("Failed to connect to Redis: %s; unique task deduplication disabled", e)
                self._deduplication_enabled = False
                self._redis_available = False


class RequestContextAndUniqueByPayloadTask(RequestContextTask, UniqueByPayloadTask):
    abstract = True
