from enum import Enum
from typing import List, Optional, Any, Callable, Awaitable, Union, AsyncIterator, ClassVar
from .decorators import WorkflowDef, 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 TransientStepError(PgqrsError): ...
class RetriesExhaustedError(PgqrsError): ...
class StepNotReadyError(PgqrsError): ...
class PausedError(PgqrsError): ...

class DurabilityMode:
    DURABLE: ClassVar["DurabilityMode"]
    LOCAL: ClassVar["DurabilityMode"]
    @property
    def value(self) -> str: ...

class WorkerStatus(Enum):
    Ready: "WorkerStatus"
    Polling: "WorkerStatus"
    Suspended: "WorkerStatus"
    Interrupted: "WorkerStatus"
    Stopped: "WorkerStatus"

class WorkflowStatus(Enum):
    Queued: "WorkflowStatus"
    Running: "WorkflowStatus"
    Paused: "WorkflowStatus"
    Success: "WorkflowStatus"
    Error: "WorkflowStatus"

class StepResultStatus(Enum):
    Execute: "StepResultStatus"
    Skipped: "StepResultStatus"

class BackoffStrategy:
    """Backoff strategy for step retries."""

    @staticmethod
    def fixed(delay_seconds: int) -> "BackoffStrategy":
        """Create a fixed delay backoff strategy."""
        ...

    @staticmethod
    def exponential(base_seconds: int, max_seconds: int) -> "BackoffStrategy":
        """Create an exponential backoff strategy: delay = base * 2^attempt."""
        ...

    @staticmethod
    def exponential_with_jitter(
        base_seconds: int, max_seconds: int
    ) -> "BackoffStrategy":
        """Create an exponential backoff with jitter strategy (±25%)."""
        ...

class StepRetryPolicy:
    """
    Retry policy for workflow steps.

    Configures automatic retry behavior when steps fail with transient errors.

    Example:
        >>> policy = StepRetryPolicy(
        ...     max_attempts=5,
        ...     backoff=BackoffStrategy.exponential_with_jitter(
        ...         base_seconds=2, max_seconds=60
        ...     )
        ... )
    """

    def __init__(
        self,
        max_attempts: int = 3,
        backoff: Optional[BackoffStrategy] = None,
    ) -> None:
        """
        Create a new retry policy.

        Args:
            max_attempts: Maximum number of retry attempts (default: 3)
            backoff: Backoff strategy (default: exponential_with_jitter with base=1s, max=60s)
        """
        ...

    @property
    def max_attempts(self) -> int:
        """Maximum number of retry attempts."""
        ...

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: ...
    @property
    def max_payload_size_bytes(self) -> int: ...
    @max_payload_size_bytes.setter
    def max_payload_size_bytes(self, value: int) -> None: ...
    @property
    def max_string_length(self) -> int: ...
    @max_string_length.setter
    def max_string_length(self, value: int) -> None: ...
    @property
    def max_object_depth(self) -> int: ...
    @max_object_depth.setter
    def max_object_depth(self, value: int) -> None: ...
    @property
    def forbidden_keys(self) -> List[str]: ...
    @forbidden_keys.setter
    def forbidden_keys(self, value: List[str]) -> None: ...
    @property
    def required_keys(self) -> List[str]: ...
    @required_keys.setter
    def required_keys(self, value: List[str]) -> None: ...
    @property
    def max_enqueue_per_second(self) -> Optional[int]: ...
    @max_enqueue_per_second.setter
    def max_enqueue_per_second(self, value: Optional[int]) -> None: ...
    @property
    def max_enqueue_burst(self) -> Optional[int]: ...
    @max_enqueue_burst.setter
    def max_enqueue_burst(self, value: Optional[int]) -> None: ...
    @property
    def s3_mode(self) -> DurabilityMode: ...
    @s3_mode.setter
    def s3_mode(self, value: DurabilityMode) -> 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": ...
    async def queue(self, name: str) -> "QueueInfo": ...
    async def get_workers(self) -> "Workers": ...
    async def get_queues(self) -> "Queues": ...
    async def get_messages(self) -> "Messages": ...
    async def get_workflows(self) -> "Workflows": ...
    async def get_workflow_runs(self) -> "WorkflowRuns": ...
    async def get_workflow_steps(self) -> "WorkflowSteps": ...
    async def bootstrap(self) -> None: ...

class S3StoreHandle:
    async def snapshot(self) -> None: ...
    async def sync(self) -> None: ...

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

class Consumer:
    @property
    def worker_id(self) -> int: ...
    async def status(self) -> WorkerStatus: ...
    async def dequeue(
        self, batch_size: Optional[int] = None
    ) -> List["QueueMessage"]: ...
    async def dequeue_many_with_delay(
        self, limit: int, vt: int
    ) -> List["QueueMessage"]: ...
    async def dequeue_at(
        self, limit: int, vt: int, current_time: str
    ) -> List["QueueMessage"]: ...
    async def extend_vt(self, message_id: int, seconds: int) -> bool: ...
    async def archive(self, message_id: int) -> bool: ...
    async def archive_many(self, message_ids: List[int]) -> List[bool]: ...
    async def delete(self, message_id: int) -> bool: ...
    async def delete_many(self, message_ids: List[int]) -> List[bool]: ...
    async def release_with_visibility(
        self, message_id: int, visible_at: str
    ) -> bool: ...
    async def release_messages(self, message_ids: List[int]) -> int: ...

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]: ...
    @property
    def archived_at(self) -> Optional[str]: ...

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) -> WorkerStatus: ...
    @property
    def queue_id(self) -> Optional[int]: ...
    @property
    def heartbeat_at(self) -> str: ...

class WorkflowRecord:
    @property
    def id(self) -> int: ...
    @property
    def name(self) -> str: ...
    @property
    def queue_id(self) -> int: ...

class RunRecord:
    @property
    def id(self) -> int: ...
    @property
    def workflow_id(self) -> int: ...
    @property
    def status(self) -> WorkflowStatus: ...
    @property
    def input(self) -> Optional[Any]: ...
    @property
    def output(self) -> Optional[Any]: ...
    @property
    def error(self) -> Optional[Any]: ...
    @property
    def created_at(self) -> str: ...
    @property
    def updated_at(self) -> str: ...

class StepRecord:
    @property
    def id(self) -> int: ...
    @property
    def run_id(self) -> int: ...
    @property
    def step_name(self) -> str: ...
    @property
    def status(self) -> WorkflowStatus: ...
    @property
    def input(self) -> Optional[Any]: ...
    @property
    def output(self) -> Optional[Any]: ...
    @property
    def error(self) -> Optional[Any]: ...
    @property
    def retry_at(self) -> Optional[str]: ...
    @property
    def created_at(self) -> str: ...
    @property
    def updated_at(self) -> str: ...

class Workflows:
    async def get_by_name(self, name: str) -> WorkflowRecord: ...

class WorkflowRuns:
    async def list(self) -> List[RunRecord]: ...

class WorkflowSteps:
    async def list(self) -> List[StepRecord]: ...

class Workers:
    async def count(self) -> int: ...
    async def list(self) -> List[WorkerInfo]: ...
    async def get(self, worker_id: int) -> 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 list_archived_by_queue(self, queue_id: int) -> List[QueueMessage]: ...

class Admin:
    def __init__(self, store: Store) -> None: ...
    async def install(self) -> None: ...
    async def verify(self) -> None: ...
    async def delete_queue(self, name: str) -> bool: ...
    async def release_worker_messages(self, worker_id: int) -> int: ...
    async def create_workflow(self, name: str, arg: Any) -> "Run": ...

class Run:
    def id(self) -> int: ...
    @property
    def current_time(self) -> Optional[str]: ...
    def with_time(self, time: str) -> "Run": ...
    async def start(self) -> None: ...
    async def fail(self, error: str) -> None: ...
    async def success(self, result: Any) -> None: ...
    async def complete(self, result: Any) -> None: ...
    async def refresh(self) -> "Run": ...
    async def pause(self, message: str, resume_after_seconds: int) -> None: ...
    async def acquire_step(
        self, step_name: str, current_time: Optional[str] = None
    ) -> "StepResult":
        """
        Acquire a step for execution.

        Args:
            step_name: Unique identifier for the step
            current_time: Optional ISO 8601 timestamp for deterministic testing (e.g., "2024-01-15T10:30:00Z")

        Returns:
            StepResult with status and optional guard

        Example:
            # Normal usage (uses system time)
            step_res = await workflow.acquire_step("my_step")

            # Testing with controlled time
            step_res = await workflow.acquire_step("my_step", current_time="2024-01-15T10:30:00Z")
        """
        ...

class StepResult:
    @property
    def status(self) -> StepResultStatus: ...
    @property
    def value(self) -> Any: ...
    @property
    def guard(self) -> Optional["StepGuard"]: ...

class StepGuard:
    async def success(self, result: Any) -> None: ...
    async def fail(self, error: str) -> None: ...
    async def fail_transient(
        self, code: str, message: str, retry_after: Optional[float] = None
    ) -> None:
        """
        Fail the step with a transient error that triggers automatic retry.

        Args:
            code: Error code for classification (e.g., "TIMEOUT", "RATE_LIMITED")
            message: Human-readable error message
            retry_after: Optional custom delay in seconds before retry (e.g., from Retry-After header)

        Example:
            await guard.fail_transient("TIMEOUT", "Connection timeout")
            await guard.fail_transient("RATE_LIMITED", "Too many requests", retry_after=60.0)
        """
        ...

async def connect(dsn: str) -> Store: ...
async def connect_with(config: Config) -> Store: ...
def as_s3(store: Store) -> S3StoreHandle: ...
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: ...

def workflow() -> "WorkflowBuilder": ...
def run() -> "RunBuilder": ...
def step() -> "StepBuilder": ...

class WorkflowBuilder:
    def name(self, name: str) -> "WorkflowBuilder": ...
    def store(self, store: Store) -> "WorkflowBuilder": ...
    def id(self, id: int) -> "WorkflowBuilder": ...
    async def create(self) -> WorkflowRecord: ...
    def trigger(self, input: Any) -> "WorkflowTriggerBuilder": ...

class WorkflowTriggerBuilder:
    async def execute(self) -> QueueMessage: ...

class RunBuilder:
    def message(self, message: QueueMessage) -> "RunBuilder": ...
    def store(self, store: Store) -> "RunBuilder": ...
    async def execute(self) -> Run: ...
    async def result(self) -> Any: ...
    async def get(self) -> Optional[Any]: ...

class StepBuilder:
    def run(self, run: Run) -> "StepBuilder": ...
    def name(self, name: str) -> "StepBuilder": ...
    def id(self, id: str) -> "StepBuilder": ...
    def with_time(self, time: str) -> "StepBuilder": ...
    async def execute(self) -> "StepResult": ...
