from typing import List, Optional, Any, Callable, Awaitable, Union, AsyncIterator
from .decorators import workflow as workflow, step as step

class PgqrsError(Exception): ...
class PgqrsConnectionError(PgqrsError): ...
class QueueNotFoundError(PgqrsError): ...
class WorkerNotFoundError(PgqrsError): ...
class QueueAlreadyExistsError(PgqrsError): ...
class MessageNotFoundError(PgqrsError): ...
class SerializationError(PgqrsError): ...
class ConfigError(PgqrsError): ...
class RateLimitedError(PgqrsError): ...
class ValidationError(PgqrsError): ...
class TimeoutError(PgqrsError): ...
class InternalError(PgqrsError): ...
class StateTransitionError(PgqrsError): ...

class Config:
    def __init__(self, dsn: str, schema: Optional[str] = None, max_connections: Optional[int] = None) -> None: ...
    @staticmethod
    def from_dsn(dsn: str) -> "Config": ...
    @property
    def dsn(self) -> str: ...
    @property
    def schema(self) -> str: ...
    @schema.setter
    def schema(self, value: str) -> None: ...
    @property
    def max_connections(self) -> int: ...
    @max_connections.setter
    def max_connections(self, value: int) -> None: ...
    @property
    def connection_timeout_seconds(self) -> int: ...
    @connection_timeout_seconds.setter
    def connection_timeout_seconds(self, value: int) -> None: ...
    @property
    def default_lock_time_seconds(self) -> int: ...
    @default_lock_time_seconds.setter
    def default_lock_time_seconds(self, value: int) -> None: ...
    @property
    def default_max_batch_size(self) -> int: ...
    @default_max_batch_size.setter
    def default_max_batch_size(self, value: int) -> None: ...
    @property
    def max_read_ct(self) -> int: ...
    @max_read_ct.setter
    def max_read_ct(self, value: int) -> None: ...
    @property
    def heartbeat_interval_seconds(self) -> int: ...
    @heartbeat_interval_seconds.setter
    def heartbeat_interval_seconds(self, value: int) -> None: ...

class Store:
    async def producer(self, queue: str) -> "Producer": ...
    async def consumer(self, queue: str) -> "Consumer": ...
    def consume_iter(self, queue: str, poll_interval_ms: Optional[int] = None) -> "ConsumerIterator": ...

class Producer:
    async def enqueue(self, payload: Any) -> int: ...
    async def enqueue_delayed(self, payload: Any, delay_seconds: int) -> int: ...

class Consumer:
    async def dequeue(self, batch_size: Optional[int] = None) -> List["QueueMessage"]: ...
    async def extend_vt(self, message_id: int, seconds: int) -> bool: ...
    async def archive(self, message_id: int) -> bool: ...
    async def delete(self, message_id: int) -> bool: ...

class ConsumerIterator:
    def __aiter__(self) -> "ConsumerIterator": ...
    async def __anext__(self) -> "QueueMessage": ...
    async def archive(self, message_id: int) -> bool: ...
    async def delete(self, message_id: int) -> bool: ...
    async def close(self) -> None: ...
    async def __aenter__(self) -> "ConsumerIterator": ...
    async def __aexit__(
        self,
        exc_type: Optional[Any],
        exc_value: Optional[Any],
        traceback: Optional[Any],
    ) -> None: ...

class QueueMessage:
    @property
    def id(self) -> int: ...
    @property
    def queue_id(self) -> int: ...
    @property
    def payload(self) -> Any: ...
    @property
    def vt(self) -> str: ...
    @property
    def enqueued_at(self) -> str: ...
    @property
    def read_ct(self) -> int: ...
    @property
    def dequeued_at(self) -> Optional[str]: ...
    @property
    def producer_worker_id(self) -> Optional[int]: ...
    @property
    def consumer_worker_id(self) -> Optional[int]: ...

class QueueInfo:
    @property
    def id(self) -> int: ...
    @property
    def queue_name(self) -> str: ...

class WorkerInfo:
    @property
    def id(self) -> int: ...
    @property
    def hostname(self) -> str: ...
    @property
    def status(self) -> str: ...

class ArchivedMessage:
    @property
    def id(self) -> int: ...
    @property
    def queue_id(self) -> int: ...
    @property
    def original_msg_id(self) -> int: ...
    @property
    def payload(self) -> Any: ...
    @property
    def producer_worker_id(self) -> Optional[int]: ...
    @property
    def consumer_worker_id(self) -> Optional[int]: ...
    @property
    def vt(self) -> str: ...
    @property
    def dequeued_at(self) -> Optional[str]: ...
    @property
    def archived_at(self) -> str: ...
    @property
    def enqueued_at(self) -> str: ...
    @property
    def read_ct(self) -> int: ...

class Workers:
    async def count(self) -> int: ...
    async def list(self) -> List[WorkerInfo]: ...
    async def health_stats(self) -> List[Any]: ...
    async def purge_stale(self, timeout_seconds: int) -> int: ...

class Queues:
    async def count(self) -> int: ...
    async def list(self) -> List[QueueInfo]: ...
    async def purge(self, queue_name: str) -> int: ...
    async def metrics(self, queue_name: str) -> Any: ...
    async def list_metrics(self) -> List[Any]: ...

class Messages:
    async def count(self) -> int: ...
    async def count_active_by_worker(self, worker_id: int) -> int: ...
    async def move_to_dlq(self, message_id: int) -> bool: ...

class Archive:
    async def count(self) -> int: ...
    async def list_by_worker(self, worker_id: int, limit: int, offset: int) -> List[ArchivedMessage]: ...
    async def count_by_worker(self, worker_id: int) -> int: ...
    async def get(self, id: int) -> ArchivedMessage: ...
    async def delete(self, id: int) -> bool: ...
    async def dlq_count(self, max_attempts: int) -> int: ...
    async def replay(self, id: int) -> int: ...

class Admin:
    def __init__(self, store: Store) -> None: ...
    async def install(self) -> None: ...
    async def verify(self) -> None: ...
    async def create_queue(self, name: str) -> QueueInfo: ...
    async def delete_queue(self, name: str) -> bool: ...
    async def get_workers(self) -> Workers: ...
    async def get_queues(self) -> Queues: ...
    async def get_messages(self) -> Messages: ...
    async def get_archive(self) -> Archive: ...
    async def create_workflow(self, name: str, arg: Any) -> "PyWorkflow": ...

class PyWorkflow:
    async def start(self) -> None: ...
    async def fail(self, error: str) -> None: ...
    async def success(self, result: Any) -> None: ...
    async def acquire_step(self, step_id: str) -> "PyStepResult": ...

class PyStepResult:
    @property
    def status(self) -> str: ...
    @property
    def value(self) -> Any: ...
    @property
    def guard(self) -> Optional["PyStepGuard"]: ...

class PyStepGuard:
    async def success(self, result: Any) -> None: ...
    async def fail(self, error: str) -> None: ...

async def connect(dsn: str) -> Store: ...
async def connect_with(config: Config) -> Store: ...
def admin(store: Store) -> Admin: ...
async def produce(store: Store, queue: str, payload: Any) -> int: ...
async def produce_batch(store: Store, queue: str, payloads: List[Any]) -> List[int]: ...
async def consume(store: Store, queue: str, handler: Callable[[QueueMessage], Awaitable[Any]]) -> None: ...
async def consume_batch(store: Store, queue: str, batch_size: int, handler: Callable[[List[QueueMessage]], Awaitable[Any]]) -> None: ...
async def enqueue(producer: Producer, payload: Any) -> int: ...
async def enqueue_batch(producer: Producer, payloads: List[Any]) -> List[int]: ...
async def enqueue_delayed(producer: Producer, payload: Any, delay_seconds: int) -> int: ...
async def dequeue(consumer: Consumer, batch_size: int) -> List[QueueMessage]: ...
async def archive(consumer: Consumer, message: QueueMessage) -> bool: ...
async def archive_batch(consumer: Consumer, messages: List[QueueMessage]) -> bool: ...
async def delete(consumer: Consumer, message: QueueMessage) -> bool: ...
async def extend_vt(consumer: Consumer, message: QueueMessage, seconds: int) -> bool: ...
