"""Celery transport adapter for the Sandai operator client."""

from __future__ import annotations

import time
from typing import Any, Dict, Iterable, Optional

from celery.result import AsyncResult

from .config import ClientConfig
from .debug_output import PrintLogger
from .shared_transport import (
    acquire_shared_app,
    acquire_shared_backend,
    acquire_shared_broker,
    release_shared_app,
    release_shared_backend,
    release_shared_broker,
)

_PENDING_STATES = {"PENDING", "RECEIVED", "STARTED", "RETRY"}


class BoundedRedisAsyncResult(AsyncResult):
    """AsyncResult wrapper that keeps Redis result access on the bounded path."""

    def __init__(self, task_id: str, *, app: Any, transport: "CeleryTransport"):
        super().__init__(task_id, app=app)
        self.id = task_id
        self._sandai_transport = transport

    def get(self, timeout: Optional[int] = None, propagate: bool = True, **kwargs: Any) -> Any:
        effective_timeout = self._sandai_transport.config.timeout if timeout is None else int(timeout)
        result = self._sandai_transport.get_result(self, timeout=effective_timeout)
        if propagate and isinstance(result, BaseException):
            raise result
        return result

    def forget(self) -> None:
        self._sandai_transport.forget(self)

    @property
    def state(self) -> str:
        return self._sandai_transport.get_state(self)

    def ready(self) -> bool:
        return self.state not in _PENDING_STATES


class CeleryTransport:
    """Wrap shared Celery app creation and bounded task/result interactions."""

    def __init__(self, config: ClientConfig, logger: Optional[Any] = None):
        self.config = config
        self.logger = logger or PrintLogger(__name__, enabled=config.debug)
        self._broker_state = acquire_shared_broker(config)
        self._backend_state = acquire_shared_backend(config)
        self._prefer_backend_polling = self._supports_backend_polling()
        self._app_state = None
        if not self._prefer_backend_polling:
            self._app_state = acquire_shared_app(config)
        self.app = self._broker_state.app if self._prefer_backend_polling else self._app_state.app
        self._closed = False
        if self.config.debug:
            self.logger.debug(
                "transport_initialized operator=%s version=%s broker=%s backend=%s broker_pool=%s backend_pool=%s backend_ops=%s prefer_backend_polling=%s",
                self.config.operator_name,
                self.config.operator_version,
                self.config.broker_url,
                self.config.result_backend,
                self.config.broker_pool_limit,
                self.config.backend_pool_limit,
                self.config.backend_poll_concurrency,
                self._prefer_backend_polling,
            )

    def _detailed_debug_enabled(self) -> bool:
        return bool(self.config.debug and self.config.detailed_debug)

    def _log_pool_snapshot(self, *, reason: str, pool: str) -> None:
        if not self._detailed_debug_enabled():
            return
        limiter = self._broker_state.limiter if pool == "broker" else self._backend_state.limiter
        snapshot = limiter.snapshot()
        self.logger.debug(
            "pool_snapshot reason=%s pool=%s limit=%s in_use=%s available=%s queue=%s",
            reason,
            snapshot["pool"],
            snapshot["limit"],
            snapshot["in_use"],
            snapshot["available"],
            snapshot["queue"],
        )

    def _supports_backend_polling(self) -> bool:
        return self._backend_state.redis_client is not None

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self.config.debug:
            self.logger.debug(
                "transport_close operator=%s version=%s",
                self.config.operator_name,
                self.config.operator_version,
            )
        if self._app_state is not None:
            release_shared_app(self._app_state, debug=self.config.debug)
        release_shared_broker(self._broker_state, debug=self.config.debug)
        release_shared_backend(self._backend_state, debug=self.config.debug)

    def _wrap_async_result(self, async_result: AsyncResult) -> AsyncResult:
        task_id = getattr(async_result, "id", None)
        if not self._prefer_backend_polling or not task_id:
            return async_result
        return BoundedRedisAsyncResult(task_id, app=self.app, transport=self)

    def send_task(self, task_name: str, *, kwargs: Dict[str, Any], queue: str) -> AsyncResult:
        if self.config.debug:
            self.logger.debug("sending task to queue=%s payload=%s", queue, kwargs)
            self._log_pool_snapshot(reason="send_task_before_acquire", pool="broker")
        try:
            with self._broker_state.limiter.acquire(self.config.pool_acquire_timeout, debug=self.config.debug):
                connection_factory = getattr(self.app, "connection_for_write", None)
                if callable(connection_factory):
                    connection_ctx = connection_factory()
                    if hasattr(connection_ctx, "__enter__") and hasattr(connection_ctx, "__exit__"):
                        with connection_ctx as connection:
                            try:
                                from kombu import Producer

                                producer = Producer(connection, auto_declare=False)
                            except ImportError:  # pragma: no cover - test stubs without kombu
                                producer = None
                            result = self.app.send_task(
                                task_name,
                                kwargs=kwargs,
                                queue=queue,
                                connection=connection,
                                producer=producer,
                            )
                    else:
                        result = self.app.send_task(task_name, kwargs=kwargs, queue=queue)
                else:
                    result = self.app.send_task(task_name, kwargs=kwargs, queue=queue)
        except Exception:
            self._log_pool_snapshot(reason="send_task_acquire_failed", pool="broker")
            raise
        wrapped = self._wrap_async_result(result)
        if self.config.debug:
            self.logger.debug("task submitted id=%s", getattr(wrapped, "id", None))
            self._log_pool_snapshot(reason="send_task_after_submit", pool="broker")
        return wrapped

    def _poll_result_meta(self, task_id: str, *, timeout: int) -> Any:
        deadline = time.monotonic() + max(0, timeout)
        if self.config.debug:
            self.logger.debug("result_poll_start task_id=%s timeout=%s", task_id, timeout)
        while True:
            meta = self._read_backend_meta(task_id, acquire_timeout=self.config.pool_acquire_timeout)

            status = str(meta.get("status", "PENDING")) if isinstance(meta, dict) else "PENDING"
            if status not in _PENDING_STATES:
                if self.config.debug:
                    self.logger.debug("result_poll_ready task_id=%s status=%s", task_id, status)
                result = meta.get("result") if isinstance(meta, dict) else None
                if isinstance(result, BaseException):
                    raise result
                return result

            if time.monotonic() >= deadline:
                if self.config.debug:
                    self.logger.debug("result_poll_timeout task_id=%s timeout=%s", task_id, timeout)
                raise TimeoutError(f"operation timed out after {timeout} seconds")
            time.sleep(0.05)

    def _read_backend_meta(self, task_id: str, *, acquire_timeout: float) -> Dict[str, Any]:
        backend = self._backend_state.app.backend
        redis_client = self._backend_state.redis_client
        if self.config.debug:
            self.logger.debug("backend_meta_fetch task_id=%s acquire_timeout=%s", task_id, acquire_timeout)
            self._log_pool_snapshot(reason="backend_meta_fetch_before_acquire", pool="backend")
        try:
            with self._backend_state.limiter.acquire(acquire_timeout, debug=self.config.debug):
                payload = redis_client.get(backend.get_key_for_task(task_id))
        except Exception:
            self._log_pool_snapshot(reason="backend_meta_fetch_acquire_failed", pool="backend")
            raise
        if not payload:
            if self.config.debug:
                self.logger.debug("backend_meta_pending task_id=%s", task_id)
            return {"status": "PENDING", "result": None}
        if self.config.debug:
            self.logger.debug("backend_meta_hit task_id=%s payload_bytes=%s", task_id, len(payload))
        return backend.decode_result(payload)

    def collect_ready_results(
        self,
        async_results: Iterable[AsyncResult],
        *,
        wait_for_slot: bool = True,
    ) -> Dict[str, Any]:
        task_results = [
            async_result
            for async_result in async_results
            if getattr(async_result, "id", None)
        ]
        if not task_results:
            return {}

        if self.config.debug:
            self.logger.debug(
                "collect_ready_results_start count=%s wait_for_slot=%s prefer_backend_polling=%s",
                len(task_results),
                wait_for_slot,
                self._prefer_backend_polling,
            )

        if getattr(self, "_prefer_backend_polling", False):
            backend = self._backend_state.app.backend
            redis_client = self._backend_state.redis_client
            task_ids = [str(async_result.id) for async_result in task_results]
            keys = [backend.get_key_for_task(task_id) for task_id in task_ids]
            acquire_timeout = self.config.pool_acquire_timeout if wait_for_slot else 0.0
            if self.config.debug and self._detailed_debug_enabled():
                self._log_pool_snapshot(reason="collect_ready_results_before_acquire", pool="backend")
            try:
                with self._backend_state.limiter.acquire(acquire_timeout, debug=self.config.debug):
                    payloads = redis_client.mget(keys)
            except Exception:
                self._log_pool_snapshot(reason="collect_ready_results_acquire_failed", pool="backend")
                raise

            ready: Dict[str, Any] = {}
            for task_id, payload in zip(task_ids, payloads):
                meta = {"status": "PENDING", "result": None} if not payload else backend.decode_result(payload)
                status = str(meta.get("status", "PENDING")) if isinstance(meta, dict) else "PENDING"
                if status in _PENDING_STATES:
                    continue
                ready[task_id] = meta.get("result") if isinstance(meta, dict) else None
            if self.config.debug:
                self.logger.debug(
                    "collect_ready_results_done count=%s ready=%s wait_for_slot=%s",
                    len(task_results),
                    len(ready),
                    wait_for_slot,
                )
            return ready

        ready: Dict[str, Any] = {}
        for async_result in task_results:
            if getattr(async_result, "state", "PENDING") in _PENDING_STATES:
                continue
            result = async_result.get(timeout=0, propagate=False)
            ready[str(async_result.id)] = result
        if self.config.debug:
            self.logger.debug(
                "collect_ready_results_done count=%s ready=%s wait_for_slot=%s",
                len(task_results),
                len(ready),
                wait_for_slot,
            )
        return ready

    def get_result(self, async_result: AsyncResult, *, timeout: int) -> Any:
        task_id = getattr(async_result, "id", None)
        if getattr(self, "_prefer_backend_polling", False) and task_id:
            return self._poll_result_meta(task_id, timeout=timeout)
        result = async_result.get(timeout=timeout, propagate=False)
        if isinstance(result, BaseException):
            raise result
        return result

    def forget(self, async_result: AsyncResult) -> None:
        task_id = getattr(async_result, "id", None)
        backend = getattr(self._backend_state.app, "backend", None)
        if getattr(self, "_prefer_backend_polling", False) and task_id:
            if self.config.debug:
                self.logger.debug("forget_result task_id=%s via=redis", task_id)
                self._log_pool_snapshot(reason="forget_before_acquire", pool="backend")
            try:
                with self._backend_state.limiter.acquire(self.config.pool_acquire_timeout, debug=self.config.debug):
                    self._backend_state.redis_client.delete(backend.get_key_for_task(task_id))
            except Exception:
                self._log_pool_snapshot(reason="forget_acquire_failed", pool="backend")
                raise
            return
        if self.config.debug:
            self.logger.debug("forget_result task_id=%s via=celery", task_id)
        async_result.forget()

    def result_exists(self, async_result: AsyncResult) -> bool:
        task_id = getattr(async_result, "id", None)
        if getattr(self, "_prefer_backend_polling", False) and task_id:
            meta = self._read_backend_meta(task_id, acquire_timeout=self.config.pool_acquire_timeout)
            status = str(meta.get("status", "PENDING")) if isinstance(meta, dict) else "PENDING"
            return status not in _PENDING_STATES
        return getattr(async_result, "state", "PENDING") != "PENDING"

    def get_state(self, async_result: AsyncResult) -> str:
        task_id = getattr(async_result, "id", None)
        backend = getattr(self._backend_state.app, "backend", None)
        if getattr(self, "_prefer_backend_polling", False) and task_id:
            if self.config.debug and self._detailed_debug_enabled():
                self._log_pool_snapshot(reason="get_state_before_acquire", pool="backend")
            try:
                with self._backend_state.limiter.acquire(self.config.pool_acquire_timeout, debug=self.config.debug):
                    payload = self._backend_state.redis_client.get(backend.get_key_for_task(task_id))
                    meta = {"status": "PENDING", "result": None} if not payload else backend.decode_result(payload)
            except Exception:
                self._log_pool_snapshot(reason="get_state_acquire_failed", pool="backend")
                raise
            return str(meta.get("status", "PENDING")) if isinstance(meta, dict) else "PENDING"
        return str(getattr(async_result, "state", "PENDING"))
