from collections.abc import Callable
from typing import Any

from celery.backends.base import Backend

__all__ = ("AsyncBackendMixin", "BaseResultConsumer", "Drainer", "register_drainer")

def register_drainer(name: str) -> Callable[[type[Drainer]], type[Drainer]]: ...

class Drainer:
    def __init__(self, result_consumer: BaseResultConsumer) -> None: ...
    def start(self) -> None: ...
    def stop(self) -> None: ...
    def drain_events_until(
        self,
        p: Callable[[], bool],
        timeout: float | None = None,
        interval: float = 1,
        on_interval: Callable[[], None] | None = None,
        wait: Callable[[float], None] | None = None,
    ) -> None: ...
    def wait_for(
        self,
        p: Callable[[], bool],
        wait: Callable[[float], None],
        timeout: float | None = None,
    ) -> None: ...

class AsyncBackendMixin:
    @property
    def is_async(self) -> bool: ...
    def wait_for_pending(
        self,
        result: Any,
        callback: Callable[..., Any] | None = None,
        propagate: bool = True,
        **kwargs: Any,
    ) -> Any: ...
    def add_pending_result(
        self, result: Any, weak: bool = False, start_drainer: bool = True
    ) -> Any: ...
    def add_pending_results(self, results: list[Any], weak: bool = False) -> None: ...
    def remove_pending_result(self, result: Any) -> None: ...
    def on_result_fulfilled(self, result: Any) -> None: ...
    def iter_native(self, result: Any, no_ack: bool = True, **kwargs: Any) -> Any: ...

class BaseResultConsumer:
    def __init__(
        self,
        backend: Backend,
        app: Any,
        accept: Any,
        pending_results: Any,
        pending_messages: Any,
    ) -> None: ...
    def start(self, initial_task_id: str, **kwargs: Any) -> None: ...
    def stop(self) -> None: ...
    def drain_events(self, timeout: float | None = None) -> None: ...
    def drain_events_until(
        self,
        p: Callable[[], bool],
        timeout: float | None = None,
        on_interval: Callable[[], None] | None = None,
    ) -> None: ...
    def consume_from(self, task_id: str) -> None: ...
    def cancel_for(self, task_id: str) -> None: ...
    def on_state_change(self, meta: dict[str, Any], message: Any) -> None: ...
    def on_after_fork(self) -> None: ...
    def on_out_of_band_result(self, message: Any) -> None: ...
    def on_wait_for_pending(
        self, result: Any, timeout: float | None = None, **kwargs: Any
    ) -> None: ...
