"""
This type stub file was generated by pyright.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Callable, Literal, Optional, Protocol, Union, runtime_checkable
from typing_extensions import TypeAlias
from temporalio.common import WorkerDeploymentVersion

_DEFAULT_RESOURCE_ACTIVITY_MAX = ...
logger = ...
@dataclass(frozen=True)
class FixedSizeSlotSupplier:
    """A fixed-size slot supplier that will never issue more than a fixed number of slots."""
    num_slots: int
    ...


@dataclass(frozen=True)
class ResourceBasedTunerConfig:
    """Options for a :py:class:`ResourceBasedTuner` or a :py:class:`ResourceBasedSlotSupplier`.

    .. warning::
        The resource based tuner is currently experimental.
    """
    target_memory_usage: float
    target_cpu_usage: float
    ...


@dataclass(frozen=True)
class ResourceBasedSlotConfig:
    """Options for a specific slot type being used with a :py:class:`ResourceBasedSlotSupplier`.

    .. warning::
        The resource based tuner is currently experimental.
    """
    minimum_slots: Optional[int] = ...
    maximum_slots: Optional[int] = ...
    ramp_throttle: Optional[timedelta] = ...


@dataclass(frozen=True)
class ResourceBasedSlotSupplier:
    """A slot supplier that will dynamically adjust the number of slots based on resource usage.

    .. warning::
        The resource based tuner is currently experimental.
    """
    slot_config: ResourceBasedSlotConfig
    tuner_config: ResourceBasedTunerConfig
    ...


class SlotPermit:
    """A permit to use a slot for a workflow/activity/local activity task.

    You can inherit from this class to add your own data to the permit.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    ...


class SlotReserveContext(Protocol):
    """Context for reserving a slot from a :py:class:`CustomSlotSupplier`.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    slot_type: Literal["workflow", "activity", "local-activity"]
    task_queue: str
    worker_identity: str
    worker_build_id: str
    worker_deployment_version: Optional[WorkerDeploymentVersion]
    is_sticky: bool
    ...


@runtime_checkable
class WorkflowSlotInfo(Protocol):
    """Info about a workflow task slot usage.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    workflow_type: str
    is_sticky: bool
    ...


@runtime_checkable
class ActivitySlotInfo(Protocol):
    """Info about an activity task slot usage.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    activity_type: str
    ...


@runtime_checkable
class LocalActivitySlotInfo(Protocol):
    """Info about a local activity task slot usage.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    activity_type: str
    ...


SlotInfo: TypeAlias = Union[WorkflowSlotInfo, ActivitySlotInfo, LocalActivitySlotInfo]
@dataclass(frozen=True)
class SlotMarkUsedContext(Protocol):
    """Context for marking a slot used from a :py:class:`CustomSlotSupplier`.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    slot_info: SlotInfo
    permit: SlotPermit
    ...


@dataclass(frozen=True)
class SlotReleaseContext:
    """Context for releasing a slot from a :py:class:`CustomSlotSupplier`.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    slot_info: Optional[SlotInfo]
    permit: SlotPermit
    ...


class CustomSlotSupplier(ABC):
    """This class can be implemented to provide custom slot supplier behavior.

    .. warning::
        Custom slot suppliers are currently experimental.
    """
    @abstractmethod
    async def reserve_slot(self, ctx: SlotReserveContext) -> SlotPermit:
        """This function is called before polling for new tasks. Your implementation must block until a
        slot is available then return a permit to use that slot.

        The only acceptable exception to throw is :py:class:`asyncio.CancelledError`, as invocations of this method may
        be cancelled. Any other exceptions thrown will be logged and ignored.

        It is technically possible but rare, during worker shutdown, for this method to be called and return a value,
        but the Rust Core may not have a chance to _observe_ that value. In such cases the returned permit will not be
        released. The permit will, however, be forgotten and python will garbage collect it. So if you use the same slot
        supplier over the lifetime of more than one worker and it is critically important for you to clean up some
        resources associated all permits you construct, then consider using a finalizer on your returned permits.

        Args:
            ctx: The context for slot reservation.

        Returns:
            A permit to use the slot which may be populated with your own data.
        """
        ...
    
    @abstractmethod
    def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
        """This function is called when trying to reserve slots for "eager" workflow and activity tasks.
        Eager tasks are those which are returned as a result of completing a workflow task, rather than
        from polling. Your implementation must not block, and if a slot is available, return a permit
        to use that slot.

        Args:
            ctx: The context for slot reservation.

        Returns:
            Maybe a permit to use the slot which may be populated with your own data.
        """
        ...
    
    @abstractmethod
    def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
        """This function is called once a slot is actually being used to process some task, which may be
        some time after the slot was reserved originally. For example, if there is no work for a
        worker, a number of slots equal to the number of active pollers may already be reserved, but
        none of them are being used yet. This call should be non-blocking.

        Args:
            ctx: The context for marking a slot as used.
        """
        ...
    
    @abstractmethod
    def release_slot(self, ctx: SlotReleaseContext) -> None:
        """This function is called once a permit is no longer needed. This could be because the task has
        finished, whether successfully or not, or because the slot was no longer needed (ex: the number
        of active pollers decreased). This call should be non-blocking.

        Args:
            ctx: The context for releasing a slot.
        """
        ...
    


SlotSupplier: TypeAlias = Union[FixedSizeSlotSupplier, ResourceBasedSlotSupplier, CustomSlotSupplier]
class _BridgeSlotSupplierWrapper:
    def __init__(self, supplier: CustomSlotSupplier) -> None:
        ...
    
    async def reserve_slot(self, ctx: SlotReserveContext, reserve_cb: Callable[[Any], None]) -> SlotPermit:
        ...
    
    def try_reserve_slot(self, ctx: SlotReserveContext) -> Optional[SlotPermit]:
        ...
    
    def release_slot(self, ctx: SlotReleaseContext) -> None:
        ...
    
    def mark_slot_used(self, ctx: SlotMarkUsedContext) -> None:
        ...
    


class WorkerTuner(ABC):
    """WorkerTuners allow for the dynamic customization of some aspects of worker configuration"""
    @staticmethod
    def create_resource_based(*, target_memory_usage: float, target_cpu_usage: float, workflow_config: Optional[ResourceBasedSlotConfig] = ..., activity_config: Optional[ResourceBasedSlotConfig] = ..., local_activity_config: Optional[ResourceBasedSlotConfig] = ...) -> WorkerTuner:
        """Create a resource-based tuner with the provided options."""
        ...
    
    @staticmethod
    def create_fixed(*, workflow_slots: Optional[int], activity_slots: Optional[int], local_activity_slots: Optional[int]) -> WorkerTuner:
        """Create a fixed-size tuner with the provided number of slots. Any unspecified slots will default to 100."""
        ...
    
    @staticmethod
    def create_composite(*, workflow_supplier: SlotSupplier, activity_supplier: SlotSupplier, local_activity_supplier: SlotSupplier) -> WorkerTuner:
        """Create a tuner composed of the provided slot suppliers."""
        ...
    


@dataclass(frozen=True)
class _CompositeTuner(WorkerTuner):
    """This tuner allows for different slot suppliers for different slot types."""
    workflow_slot_supplier: SlotSupplier
    activity_slot_supplier: SlotSupplier
    local_activity_slot_supplier: SlotSupplier
    ...


