"""Async transport helpers for the Sandai operator client."""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Dict, Optional, Sequence
from urllib.parse import urlparse

try:  # pragma: no cover - exercised in runtime environments with optional deps
    from kombu.serialization import loads as kombu_loads
except ImportError:  # pragma: no cover
    kombu_loads = None

try:  # pragma: no cover - exercised in runtime environments with optional deps
    from redis.asyncio import Redis
    from redis.asyncio import from_url as redis_from_url
except ImportError:  # pragma: no cover
    Redis = Any  # type: ignore[misc,assignment]
    redis_from_url = None

from .config import ClientConfig
from .transport import CeleryTransport

logger = logging.getLogger(__name__)


def _serializer_content_type(serializer: str) -> str:
    if serializer == "msgpack":
        return "application/x-msgpack"
    return "application/json"


def _serializer_content_encoding(serializer: str) -> str:
    if serializer == "msgpack":
        return "binary"
    return "utf-8"


def _accept_content_types(values: Sequence[str]) -> Sequence[str]:
    content_types = []
    for value in values:
        if value == "json":
            content_types.append("application/json")
        elif value == "msgpack":
            content_types.append("application/x-msgpack")
        else:
            content_types.append(value)
    return content_types


def _is_redis_url(url: str) -> bool:
    return urlparse(url).scheme in {"redis", "rediss"}


class AsyncResultHub:
    """Shared async Redis poller for many Celery result keys."""

    def __init__(
        self,
        redis_client: Redis,
        *,
        result_serializer: str,
        accept_content: Sequence[str],
        poll_interval: float = 0.1,
        logger_: Optional[logging.Logger] = None,
    ) -> None:
        self.redis = redis_client
        self.result_serializer = result_serializer
        self.accept_content = tuple(_accept_content_types(accept_content))
        self.poll_interval = poll_interval
        self.logger = logger_ or logger
        self._pending: Dict[str, list[asyncio.Future[Any]]] = {}
        self._lock = asyncio.Lock()
        self._wakeup = asyncio.Event()
        self._closed = False
        self._poller_task: Optional[asyncio.Task[None]] = None

    def _result_key(self, task_id: str) -> str:
        return f"celery-task-meta-{task_id}"

    async def start(self) -> None:
        if self._poller_task is None:
            self._poller_task = asyncio.create_task(self._poll_loop())

    async def close(self) -> None:
        self._closed = True
        self._wakeup.set()
        if self._poller_task is not None:
            await self._poller_task
            self._poller_task = None

    async def wait_for_result(self, task_id: str, *, timeout: Optional[float]) -> Any:
        await self.start()
        loop = asyncio.get_running_loop()
        future: asyncio.Future[Any] = loop.create_future()
        async with self._lock:
            self._pending.setdefault(task_id, []).append(future)
            self._wakeup.set()

        try:
            if timeout is None:
                return await future
            return await asyncio.wait_for(future, timeout=timeout)
        finally:
            if future.cancelled():
                async with self._lock:
                    futures = self._pending.get(task_id, [])
                    if future in futures:
                        futures.remove(future)
                        if not futures:
                            self._pending.pop(task_id, None)

    async def result_exists(self, task_id: str) -> bool:
        return bool(await self.redis.exists(self._result_key(task_id)))

    async def get_state(self, task_id: str) -> str:
        value = await self.redis.get(self._result_key(task_id))
        if value is None:
            return "PENDING"
        meta = self._decode_meta(value)
        return str(meta.get("status", "PENDING"))

    async def forget(self, task_id: str) -> None:
        await self.redis.delete(self._result_key(task_id))

    def _decode_meta(self, payload: Any) -> Dict[str, Any]:
        if kombu_loads is None:
            raise RuntimeError("kombu is required for AsyncOperatorClient result decoding")
        return kombu_loads(
            payload,
            content_type=_serializer_content_type(self.result_serializer),
            content_encoding=_serializer_content_encoding(self.result_serializer),
            accept=self.accept_content,
        )

    async def _poll_loop(self) -> None:
        while not self._closed:
            await self._wakeup.wait()
            self._wakeup.clear()

            while True:
                async with self._lock:
                    task_ids = list(self._pending.keys())
                if not task_ids:
                    break

                values = await self.redis.mget([self._result_key(task_id) for task_id in task_ids])
                ready = 0
                for task_id, value in zip(task_ids, values):
                    if value is None:
                        continue
                    ready += 1
                    meta = self._decode_meta(value)
                    async with self._lock:
                        futures = self._pending.pop(task_id, [])
                    for future in futures:
                        if not future.done():
                            future.set_result(meta.get("result"))

                if ready == 0:
                    try:
                        await asyncio.wait_for(self._wakeup.wait(), timeout=self.poll_interval)
                        self._wakeup.clear()
                    except asyncio.TimeoutError:
                        pass


class AsyncTransportAdapter:
    """Async bridge over the existing Celery transport with Redis fast-paths."""

    def __init__(self, config: ClientConfig, logger_: Optional[logging.Logger] = None):
        self.config = config
        self.logger = logger_ or logger
        self.celery = CeleryTransport(config, logger=self.logger)
        self.app = self.celery.app
        self._result_hub: Optional[AsyncResultHub] = None
        self._redis: Optional[Redis] = None
        self._closed = False

        if _is_redis_url(config.result_backend):
            if redis_from_url is None:
                raise RuntimeError("redis is required for AsyncOperatorClient Redis result polling")
            self._redis = redis_from_url(config.result_backend, decode_responses=False)
            self._result_hub = AsyncResultHub(
                self._redis,
                result_serializer=config.celery.result_serializer,
                accept_content=config.celery.accept_content or [],
                logger_=self.logger,
            )

    async def aclose(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self._result_hub is not None:
            await self._result_hub.close()
        if self._redis is not None:
            await self._redis.aclose()

    async def send_task(self, task_name: str, *, kwargs: Dict[str, Any], queue: str) -> Any:
        return await asyncio.to_thread(
            self.celery.send_task,
            task_name,
            kwargs=kwargs,
            queue=queue,
        )

    async def get_result(self, async_result: Any, *, timeout: int) -> Any:
        if self._result_hub is not None:
            return await self._result_hub.wait_for_result(async_result.id, timeout=timeout)
        return await asyncio.to_thread(self.celery.get_result, async_result, timeout=timeout)

    async def forget(self, async_result: Any) -> None:
        if self._result_hub is not None:
            await self._result_hub.forget(async_result.id)
            return
        await asyncio.to_thread(self.celery.forget, async_result)

    async def result_exists(self, async_result: Any) -> bool:
        if self._result_hub is not None:
            return await self._result_hub.result_exists(async_result.id)
        return await asyncio.to_thread(lambda: async_result.state != "PENDING")

    async def get_state(self, async_result: Any) -> str:
        if self._result_hub is not None:
            return await self._result_hub.get_state(async_result.id)
        return await asyncio.to_thread(lambda: str(async_result.state))
