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

import contextlib
import threading
from dataclasses import dataclass
from typing import Any, Literal, TYPE_CHECKING, Union
from collections.abc import Iterator
from types import FunctionType
from google.protobuf.message import Message
from streamlit.proto.Block_pb2 import Block
from streamlit.runtime.caching.cache_type import CacheType

if TYPE_CHECKING:
    ...
@dataclass(frozen=True)
class MediaMsgData:
    media: bytes | str
    mimetype: str
    media_id: str
    ...


@dataclass(frozen=True)
class ElementMsgData:
    """An element's message and related metadata for
    replaying that element's function call.

    media_data is filled in iff this is a media element (image, audio, video).
    """
    delta_type: str
    message: Message
    id_of_dg_called_on: str
    returned_dgs_id: str
    media_data: list[MediaMsgData] | None = ...


@dataclass(frozen=True)
class BlockMsgData:
    message: Block
    id_of_dg_called_on: str
    returned_dgs_id: str
    ...


MsgData = Union[ElementMsgData, BlockMsgData]
@dataclass
class CachedResult:
    """The full results of calling a cache-decorated function, enough to
    replay the st functions called while executing it.
    """
    value: Any
    messages: list[MsgData]
    main_id: str
    sidebar_id: str
    ...


class CachedMessageReplayContext(threading.local):
    """A utility for storing messages generated by `st` commands called inside
    a cached function.

    Data is stored in a thread-local object, so it's safe to use an instance
    of this class across multiple threads.
    """
    def __init__(self, cache_type: CacheType) -> None:
        ...
    
    def __repr__(self) -> str:
        ...
    
    @contextlib.contextmanager
    def calling_cached_function(self, func: FunctionType) -> Iterator[None]:
        """Context manager that should wrap the invocation of a cached function.
        It allows us to track any `st.foo` messages that are generated from inside the
        function for playback during cache retrieval.
        """
        ...
    
    def save_element_message(self, delta_type: str, element_proto: Message, invoked_dg_id: str, used_dg_id: str, returned_dg_id: str) -> None:
        """Record the element protobuf as having been produced during any currently
        executing cached functions, so they can be replayed any time the function's
        execution is skipped because they're in the cache.
        """
        ...
    
    def save_block_message(self, block_proto: Block, invoked_dg_id: str, used_dg_id: str, returned_dg_id: str) -> None:
        ...
    
    def select_dg_to_save(self, invoked_id: str, acting_on_id: str) -> str:
        """Select the id of the DG that this message should be invoked on
        during message replay.

        See Note [DeltaGenerator method invocation]

        invoked_id is the DG the st function was called on, usually `st._main`.
        acting_on_id is the DG the st function ultimately runs on, which may be different
        if the invoked DG delegated to another one because it was in a `with` block.
        """
        ...
    
    def save_image_data(self, image_data: bytes | str, mimetype: str, image_id: str) -> None:
        ...
    


def replay_cached_messages(result: CachedResult, cache_type: CacheType, cached_func: FunctionType) -> None:
    """Replay the st element function calls that happened when executing a
    cache-decorated function.

    When a cache function is executed, we record the element and block messages
    produced, and use those to reproduce the DeltaGenerator calls, so the elements
    will appear in the web app even when execution of the function is skipped
    because the result was cached.

    To make this work, for each st function call we record an identifier for the
    DG it was effectively called on (see Note [DeltaGenerator method invocation]).
    We also record the identifier for each DG returned by an st function call, if
    it returns one. Then, for each recorded message, we get the current DG instance
    corresponding to the DG the message was originally called on, and enqueue the
    message using that, recording any new DGs produced in case a later st function
    call is on one of them.
    """
    ...

def show_widget_replay_deprecation(decorator: Literal["cache_data", "cache_resource"]) -> None:
    ...

