"""Process-global shared Celery transport resources.

This module keeps broker/backend connection pools bounded at the process level
while allowing many operator clients to share the same underlying Celery apps.
"""

from __future__ import annotations

import os
import threading
import time
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Iterator, Optional, Tuple
from urllib.parse import urlparse

from celery import Celery

try:  # pragma: no cover - exercised in runtime environments with optional deps
    from redis import Redis
    from redis 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 .debug_output import PrintLogger
from .exceptions import ClientPoolTimeoutError

logger = PrintLogger(__name__, enabled=True)
CLIENT_NAME_PREFIX = "sandai-operator-client"


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


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


def _normalize_accept_content(values: Optional[Iterable[str]]) -> Tuple[str, ...]:
    if values is None:
        return ()
    return tuple(values)


def _scope_suffix(config: ClientConfig) -> Tuple[str, ...]:
    if config.shared_transport_scope == "operator":
        return (config.operator_name, config.operator_version)
    return ()


def _client_name(role: str, endpoint: str) -> str:
    return f"{CLIENT_NAME_PREFIX}:{role}:{os.getpid()}:{abs(hash(endpoint))}"


@dataclass(frozen=True)
class AppResourceKey:
    broker_url: str
    result_backend: str
    task_serializer: str
    result_serializer: str
    accept_content: Tuple[str, ...]
    broker_pool_limit: int
    backend_pool_limit: int
    scope: Tuple[str, ...] = ()


@dataclass(frozen=True)
class BrokerResourceKey:
    broker_url: str
    task_serializer: str
    accept_content: Tuple[str, ...]
    pool_limit: int
    scope: Tuple[str, ...] = ()


@dataclass(frozen=True)
class BackendResourceKey:
    result_backend: str
    result_serializer: str
    accept_content: Tuple[str, ...]
    pool_limit: int
    operation_limit: int
    scope: Tuple[str, ...] = ()


class FairLeaseLimiter:
    """A FIFO limiter so different operators do not starve each other."""

    def __init__(self, name: str, limit: int):
        self.name = name
        self.limit = limit
        self._available = limit
        self._cond = threading.Condition()
        self._queue: deque[object] = deque()

    @contextmanager
    def acquire(self, timeout: float, *, debug: bool = False) -> Iterator[None]:
        self._acquire(timeout, debug=debug)
        try:
            yield
        finally:
            self.release(debug=debug)

    def _acquire(self, timeout: float, *, debug: bool = False) -> None:
        deadline = None if timeout is None else time.monotonic() + max(0.0, timeout)
        ticket = object()
        with self._cond:
            self._queue.append(ticket)
            if debug:
                logger.debug(
                    "pool_wait_enter pool=%s timeout=%s queue=%s available=%s limit=%s",
                    self.name,
                    timeout,
                    len(self._queue),
                    self._available,
                    self.limit,
                )
            while True:
                is_front = self._queue and self._queue[0] is ticket
                if is_front and self._available > 0:
                    self._queue.popleft()
                    self._available -= 1
                    if debug:
                        logger.debug(
                            "pool_wait_acquired pool=%s queue=%s available=%s limit=%s",
                            self.name,
                            len(self._queue),
                            self._available,
                            self.limit,
                        )
                    return

                if deadline is None:
                    self._cond.wait()
                    continue

                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    try:
                        self._queue.remove(ticket)
                    except ValueError:
                        pass
                    if debug:
                        logger.debug(
                            "pool_wait_timeout pool=%s timeout=%s queue=%s available=%s limit=%s",
                            self.name,
                            timeout,
                            len(self._queue),
                            self._available,
                            self.limit,
                        )
                    raise ClientPoolTimeoutError(self.name, timeout, self.limit)
                self._cond.wait(timeout=remaining)

    def release(self, *, debug: bool = False) -> None:
        with self._cond:
            self._available += 1
            if self._available > self.limit:
                self._available = self.limit
            if debug:
                logger.debug(
                    "pool_wait_released pool=%s queue=%s available=%s limit=%s",
                    self.name,
                    len(self._queue),
                    self._available,
                    self.limit,
                )
            self._cond.notify_all()

    def snapshot(self) -> Dict[str, int | str]:
        with self._cond:
            return {
                "pool": self.name,
                "limit": self.limit,
                "available": self._available,
                "in_use": self.limit - self._available,
                "queue": len(self._queue),
            }


@dataclass
class SharedAppState:
    key: AppResourceKey
    app: Celery
    ref_count: int = 0


@dataclass
class SharedBrokerState:
    key: BrokerResourceKey
    app: Celery
    limiter: FairLeaseLimiter
    ref_count: int = 0


@dataclass
class SharedBackendState:
    key: BackendResourceKey
    app: Celery
    limiter: FairLeaseLimiter
    redis_client: Optional[Redis] = None
    ref_count: int = 0


_registry_lock = threading.Lock()
_shared_apps: Dict[AppResourceKey, SharedAppState] = {}
_shared_brokers: Dict[BrokerResourceKey, SharedBrokerState] = {}
_shared_backends: Dict[BackendResourceKey, SharedBackendState] = {}


def _make_app_key(config: ClientConfig) -> AppResourceKey:
    return AppResourceKey(
        broker_url=config.broker_url,
        result_backend=config.result_backend,
        task_serializer=config.celery.task_serializer,
        result_serializer=config.celery.result_serializer,
        accept_content=_normalize_accept_content(config.celery.accept_content),
        broker_pool_limit=config.broker_pool_limit,
        backend_pool_limit=config.backend_pool_limit,
        scope=_scope_suffix(config),
    )


def _make_broker_key(config: ClientConfig) -> BrokerResourceKey:
    return BrokerResourceKey(
        broker_url=config.broker_url,
        task_serializer=config.celery.task_serializer,
        accept_content=_normalize_accept_content(config.celery.accept_content),
        pool_limit=config.broker_pool_limit,
        scope=_scope_suffix(config),
    )


def _make_backend_key(config: ClientConfig) -> BackendResourceKey:
    return BackendResourceKey(
        result_backend=config.result_backend,
        result_serializer=config.celery.result_serializer,
        accept_content=_normalize_accept_content(config.celery.accept_content),
        pool_limit=config.backend_pool_limit,
        operation_limit=min(config.backend_pool_limit, config.backend_poll_concurrency),
        scope=_scope_suffix(config),
    )


def _create_shared_app(key: AppResourceKey) -> Celery:
    app = Celery(f"sandai_shared_transport_{abs(hash(key))}")
    broker_transport_options: Dict[str, Any] = {}
    if _is_celery_redis_url(key.broker_url):
        broker_transport_options["max_connections"] = key.broker_pool_limit
        broker_transport_options["client_name"] = _client_name("broker", key.broker_url)

    conf: Dict[str, Any] = {
        "broker_url": key.broker_url,
        "result_backend": key.result_backend,
        "task_serializer": key.task_serializer,
        "result_serializer": key.result_serializer,
        "accept_content": list(key.accept_content),
        "broker_pool_limit": key.broker_pool_limit,
        "broker_connection_retry_on_startup": True,
        "broker_connection_retry": True,
        "broker_transport_options": broker_transport_options,
        "result_backend_max_retries": 3,
    }
    if _is_celery_redis_url(key.result_backend):
        conf["redis_max_connections"] = key.backend_pool_limit
        conf["redis_client_name"] = _client_name("backend", key.result_backend)
    app.conf.update(**conf)
    return app


def _create_broker_app(key: BrokerResourceKey) -> Celery:
    app = Celery(f"sandai_shared_broker_{abs(hash(key))}")
    broker_transport_options: Dict[str, Any] = {}
    if _is_celery_redis_url(key.broker_url):
        broker_transport_options["max_connections"] = key.pool_limit
        broker_transport_options["client_name"] = _client_name("broker", key.broker_url)
    app.conf.update(
        broker_url=key.broker_url,
        task_serializer=key.task_serializer,
        accept_content=list(key.accept_content),
        broker_pool_limit=key.pool_limit,
        broker_connection_retry_on_startup=True,
        broker_connection_retry=True,
        broker_transport_options=broker_transport_options,
    )
    return app


def _create_backend_app(key: BackendResourceKey) -> Celery:
    app = Celery(f"sandai_shared_backend_{abs(hash(key))}")
    conf: Dict[str, Any] = {
        "result_backend": key.result_backend,
        "result_serializer": key.result_serializer,
        "accept_content": list(key.accept_content),
        "result_backend_max_retries": 3,
    }
    if _is_celery_redis_url(key.result_backend):
        conf["redis_max_connections"] = key.pool_limit
        conf["redis_client_name"] = _client_name("backend", key.result_backend)
    app.conf.update(**conf)
    return app


def _create_backend_redis_client(key: BackendResourceKey) -> Optional[Redis]:
    if not _is_direct_redis_url(key.result_backend) or redis_from_url is None:
        return None
    return redis_from_url(
        key.result_backend,
        decode_responses=False,
        max_connections=key.pool_limit,
        client_name=_client_name("backend", key.result_backend),
    )


def acquire_shared_app(config: ClientConfig) -> SharedAppState:
    key = _make_app_key(config)
    with _registry_lock:
        state = _shared_apps.get(key)
        if state is None:
            state = SharedAppState(key=key, app=_create_shared_app(key))
            _shared_apps[key] = state
            if config.debug:
                logger.debug("shared_app_created key=%s", key)
        state.ref_count += 1
        if config.debug:
            logger.debug("shared_app_acquired key=%s ref_count=%s", key, state.ref_count)
        return state


def acquire_shared_broker(config: ClientConfig) -> SharedBrokerState:
    key = _make_broker_key(config)
    with _registry_lock:
        state = _shared_brokers.get(key)
        if state is None:
            state = SharedBrokerState(
                key=key,
                app=_create_broker_app(key),
                limiter=FairLeaseLimiter(f"broker:{key.broker_url}", key.pool_limit),
            )
            _shared_brokers[key] = state
            if config.debug:
                logger.debug("shared_broker_created key=%s", key)
        state.ref_count += 1
        if config.debug:
            logger.debug("shared_broker_acquired key=%s ref_count=%s", key, state.ref_count)
        return state


def acquire_shared_backend(config: ClientConfig) -> SharedBackendState:
    key = _make_backend_key(config)
    with _registry_lock:
        state = _shared_backends.get(key)
        if state is None:
            state = SharedBackendState(
                key=key,
                app=_create_backend_app(key),
                limiter=FairLeaseLimiter(f"backend:{key.result_backend}", key.operation_limit),
                redis_client=_create_backend_redis_client(key),
            )
            _shared_backends[key] = state
            if config.debug:
                logger.debug(
                    "shared_backend_created key=%s redis_polling=%s",
                    key,
                    state.redis_client is not None,
                )
        state.ref_count += 1
        if config.debug:
            logger.debug("shared_backend_acquired key=%s ref_count=%s", key, state.ref_count)
        return state


def _best_effort_close_app(app: Celery, *, debug: bool = False) -> None:
    pool = getattr(app, "pool", None)
    producer_pool = getattr(app, "producer_pool", None)
    producer_connections = getattr(producer_pool, "connections", None)
    backend = getattr(app, "backend", None)
    result_consumer = getattr(backend, "result_consumer", None)
    client = getattr(backend, "client", None)
    connection_pool = getattr(client, "connection_pool", None)

    stop = getattr(result_consumer, "stop", None)
    if callable(stop):
        try:
            stop()
        except Exception:
            if debug:
                logger.debug("failed to stop celery result consumer", exc_info=True)

    force_close_all = getattr(producer_pool, "force_close_all", None)
    if callable(force_close_all):
        try:
            force_close_all(close_pool=False)
        except Exception:
            if debug:
                logger.debug("failed to close celery producer pool", exc_info=True)

    producer_connections_close = getattr(producer_connections, "force_close_all", None)
    if callable(producer_connections_close):
        try:
            producer_connections_close(close_pool=False)
        except Exception:
            if debug:
                logger.debug("failed to close celery producer connections", exc_info=True)

    pool_force_close_all = getattr(pool, "force_close_all", None)
    if callable(pool_force_close_all):
        try:
            pool_force_close_all(close_pool=False)
        except Exception:
            if debug:
                logger.debug("failed to close celery broker pool", exc_info=True)

    disconnect = getattr(connection_pool, "disconnect", None)
    if callable(disconnect):
        try:
            disconnect()
        except Exception:
            if debug:
                logger.debug("failed to disconnect celery backend pool", exc_info=True)

    close = getattr(app, "close", None)
    if callable(close):
        try:
            close()
        except Exception:
            if debug:
                logger.debug("failed to close celery app", exc_info=True)


def _best_effort_close_backend_client(client: Optional[Redis], *, debug: bool = False) -> None:
    if client is None:
        return

    close = getattr(client, "close", None)
    if callable(close):
        try:
            close()
        except Exception:
            if debug:
                logger.debug("failed to close redis backend client", exc_info=True)

    connection_pool = getattr(client, "connection_pool", None)
    disconnect = getattr(connection_pool, "disconnect", None)
    if callable(disconnect):
        try:
            disconnect()
        except Exception:
            if debug:
                logger.debug("failed to disconnect redis backend pool", exc_info=True)


def release_shared_app(state: SharedAppState, *, debug: bool = False) -> None:
    with _registry_lock:
        current = _shared_apps.get(state.key)
        if current is None:
            return
        current.ref_count -= 1
        if current.ref_count > 0:
            if debug:
                logger.debug("shared_app_released key=%s ref_count=%s", state.key, current.ref_count)
            return
        _shared_apps.pop(state.key, None)
    if debug:
        logger.debug("shared_app_closed key=%s", state.key)
    _best_effort_close_app(state.app, debug=debug)


def release_shared_broker(state: SharedBrokerState, *, debug: bool = False) -> None:
    with _registry_lock:
        current = _shared_brokers.get(state.key)
        if current is None:
            return
        current.ref_count -= 1
        if current.ref_count > 0:
            if debug:
                logger.debug("shared_broker_released key=%s ref_count=%s", state.key, current.ref_count)
            return
        _shared_brokers.pop(state.key, None)
    if debug:
        logger.debug("shared_broker_closed key=%s", state.key)
    _best_effort_close_app(state.app, debug=debug)


def release_shared_backend(state: SharedBackendState, *, debug: bool = False) -> None:
    with _registry_lock:
        current = _shared_backends.get(state.key)
        if current is None:
            return
        current.ref_count -= 1
        if current.ref_count > 0:
            if debug:
                logger.debug("shared_backend_released key=%s ref_count=%s", state.key, current.ref_count)
            return
        _shared_backends.pop(state.key, None)
    if debug:
        logger.debug("shared_backend_closed key=%s", state.key)
    _best_effort_close_backend_client(state.redis_client, debug=debug)
    _best_effort_close_app(state.app, debug=debug)


def reset_shared_transport_state() -> None:
    """Test helper to clear shared state between regression cases."""

    with _registry_lock:
        apps = list(_shared_apps.values())
        brokers = list(_shared_brokers.values())
        backends = list(_shared_backends.values())
        _shared_apps.clear()
        _shared_brokers.clear()
        _shared_backends.clear()

    for state in apps:
        _best_effort_close_app(state.app, debug=False)
    for state in brokers:
        _best_effort_close_app(state.app, debug=False)
    for state in backends:
        _best_effort_close_backend_client(state.redis_client)
        _best_effort_close_app(state.app, debug=False)
