# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Assert statements are allowed here since the app testing logic is used within unit tests:
# ruff: noqa: S101

from __future__ import annotations

import builtins
import textwrap
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterator, Sequence
from dataclasses import dataclass, field, fields, is_dataclass
from datetime import date, datetime, time, timedelta
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Final,
    Generic,
    NoReturn,
    TypeAlias,
    TypeVar,
    cast,
    overload,
)

from typing_extensions import Self

from streamlit import dataframe_util, util
from streamlit.elements.heading import HeadingProtoTag
from streamlit.elements.lib.layout_utils import SIZE_TO_REM_MAPPING, SpaceSize
from streamlit.elements.widgets.select_slider import SelectSliderSerde
from streamlit.elements.widgets.slider import SliderSerde, SliderStep
from streamlit.elements.widgets.time_widgets import (
    DateInputSerde,
    DateTimeInputSerde,
    DateWidgetReturn,
    TimeInputSerde,
    _parse_date_value,
)
from streamlit.proto.Alert_pb2 import Alert as AlertProto
from streamlit.proto.ButtonGroup_pb2 import ButtonGroup as ButtonGroupProto
from streamlit.proto.Checkbox_pb2 import Checkbox as CheckboxProto
from streamlit.proto.GapSize_pb2 import GapSize
from streamlit.proto.Markdown_pb2 import Markdown as MarkdownProto
from streamlit.proto.Slider_pb2 import Slider as SliderProto
from streamlit.proto.WidgetStates_pb2 import WidgetState, WidgetStates
from streamlit.runtime.state.common import TESTING_KEY, user_key_from_element_id
from streamlit.testing.v1.errors import AppTestError

if TYPE_CHECKING:
    from pandas import DataFrame as PandasDataframe

    from streamlit.proto.Block_pb2 import Block as BlockProto
    from streamlit.proto.Button_pb2 import Button as ButtonProto
    from streamlit.proto.ChatInput_pb2 import ChatInput as ChatInputProto
    from streamlit.proto.Code_pb2 import Code as CodeProto
    from streamlit.proto.ColorPicker_pb2 import ColorPicker as ColorPickerProto
    from streamlit.proto.Dataframe_pb2 import Dataframe as DataframeProto
    from streamlit.proto.DateInput_pb2 import DateInput as DateInputProto
    from streamlit.proto.DateTimeInput_pb2 import DateTimeInput as DateTimeInputProto
    from streamlit.proto.DownloadButton_pb2 import (
        DownloadButton as DownloadButtonProto,
    )
    from streamlit.proto.Element_pb2 import Element as ElementProto
    from streamlit.proto.Exception_pb2 import Exception as ExceptionProto
    from streamlit.proto.Feedback_pb2 import Feedback as FeedbackProto
    from streamlit.proto.FileUploader_pb2 import FileUploader as FileUploaderProto
    from streamlit.proto.ForwardMsg_pb2 import ForwardMsg
    from streamlit.proto.Heading_pb2 import Heading as HeadingProto
    from streamlit.proto.Image_pb2 import ImageList as ImageListProto
    from streamlit.proto.Json_pb2 import Json as JsonProto
    from streamlit.proto.MenuButton_pb2 import MenuButton as MenuButtonProto
    from streamlit.proto.Metric_pb2 import Metric as MetricProto
    from streamlit.proto.MultiSelect_pb2 import MultiSelect as MultiSelectProto
    from streamlit.proto.NumberInput_pb2 import NumberInput as NumberInputProto
    from streamlit.proto.Radio_pb2 import Radio as RadioProto
    from streamlit.proto.Selectbox_pb2 import Selectbox as SelectboxProto
    from streamlit.proto.Space_pb2 import Space as SpaceProto
    from streamlit.proto.Table_pb2 import Table as TableProto
    from streamlit.proto.Text_pb2 import Text as TextProto
    from streamlit.proto.TextArea_pb2 import TextArea as TextAreaProto
    from streamlit.proto.TextInput_pb2 import TextInput as TextInputProto
    from streamlit.proto.TimeInput_pb2 import TimeInput as TimeInputProto
    from streamlit.proto.Toast_pb2 import Toast as ToastProto
    from streamlit.proto.WidthConfig_pb2 import WidthConfig
    from streamlit.runtime.state.safe_session_state import SafeSessionState
    from streamlit.testing.v1.app_test import AppTest
    from streamlit.typing import ChatInputValue

T = TypeVar("T")

# Public ``get()`` names that are not the node ``type`` string.
_GET_TYPE_ALIASES: dict[str, str] = {
    "datetime_input": "date_time_input",
    "columns": "column",
    "help": "help_info",
    "tabs": "tab",
}


def _unknown_element_content(proto: Any) -> Any:
    """Best-effort payload for an unimplemented element's proto.

    Many display protos store content in ``body``, ``text``, or ``label``
    rather than ``value`` (for example ``st.html``).
    """
    fields = getattr(getattr(proto, "DESCRIPTOR", None), "fields_by_name", None)
    if fields:
        for name in ("value", "body", "text", "label"):
            if name in fields:
                return getattr(proto, name)
    return getattr(proto, "value", None)


# Inverse of SIZE_TO_REM_MAPPING so rem_width round-trips to the named size.
# Named sizes must map to distinct rem values that are exact in float32, or
# the lookup below silently misses and returns None; test_space_named_size
# covers every name.
_REM_TO_SPACE_SIZE: Final = {
    rem: cast("SpaceSize", name) for name, rem in SIZE_TO_REM_MAPPING.items()
}


def _space_size_from_width_config(width_config: WidthConfig) -> SpaceSize | None:
    """Reconstruct the ``st.space`` size from ``Element.width_config``.

    The Space proto does not store the size. ``st.space`` writes the same
    value into both the width and height configs and lets the frontend pick
    the relevant axis, so reading width alone recovers the original argument.
    """
    spec = width_config.WhichOneof("width_spec")
    if spec == "use_stretch":
        return "stretch"
    if spec == "pixel_width":
        return width_config.pixel_width
    if spec == "rem_width":
        return _REM_TO_SPACE_SIZE.get(width_config.rem_width)
    return None  # pragma: no cover - defensive


def _format_value_for_widget(format_func: Callable[[Any], str], value: Any) -> str:
    """Format a widget value into its option label.

    AppTest supports setting widget values with either the raw option or the
    formatted label. The latter matches what the production frontend sends to
    the backend, but may raise when passed directly to ``format_func`` if the
    function expects a raw option. To support formatted labels without hiding
    real ``format_func`` bugs, only string values fall back to themselves.
    """
    try:
        return format_func(value)
    except builtins.Exception:
        if isinstance(value, str):
            return value
        raise


@dataclass
class InitialValue:
    """Used to represent the initial value of a widget."""


# TODO: This class serves as a fallback option for elements that have not
# been implemented yet, as well as providing implementations of some
# trivial methods. It may have significantly reduced scope once all elements
# have been implemented.
# This class will not be sufficient implementation for most elements.
# Widgets need their own classes to translate interactions into the appropriate
# WidgetState and provide higher level interaction interfaces, and other elements
# have enough variation in how to get their values that most will need their
# own classes too.
@dataclass
class Element(ABC):
    """
    Element base class for testing.

    This class's methods and attributes are universal for all elements
    implemented in testing. For example, ``Caption``, ``Code``, ``Text``, and
    ``Title`` inherit from ``Element``. All widget classes also
    inherit from Element, but have additional methods specific to each
    widget type. See the ``AppTest`` class for the full list of supported
    elements.

    For all element classes, parameters of the original element can be obtained
    as properties. For example, ``Button.label``, ``Caption.help``, and
    ``Toast.icon``.

    """

    type: str = field(repr=False)
    proto: Any = field(repr=False)
    root: ElementTree = field(repr=False)
    key: str | None

    @abstractmethod
    def __init__(self, proto: Any, root: ElementTree) -> None:
        # Shared proto/root assignment for subclasses that call super().__init__.
        self.proto = proto
        self.root = root

    def __iter__(self) -> Iterator[Self]:
        yield self

    @property
    @abstractmethod
    def value(self) -> Any:
        """The value or contents of the element."""
        ...

    def __getattr__(self, name: str) -> Any:
        """Look up missing names on the element's proto, except AppTest interactions.

        ``set_value`` and ``click`` are unsupported AppTest interactions even when a
        proto defines those names as fields (pagination's ``set_value`` is a bool).
        Other missing names still fall through to the proto.
        """
        if name in {"set_value", "click"}:

            def unsupported_interaction(*_args: Any, **_kwargs: Any) -> None:
                self._raise_unsupported_interaction(name)

            return unsupported_interaction
        return getattr(self.proto, name)

    def _raise_unsupported_interaction(self, method: str) -> NoReturn:
        """Raise AppTestError: typed widgets already have interaction methods; other nodes do not."""
        key_part = f" (key={self.key!r})" if self.key else ""
        if isinstance(self, Widget):
            raise AppTestError(
                f"{method}() is not supported for {self.type}{key_part}. "
                "Use set_value() or one of this widget's typed interaction "
                "methods."
            )
        raise AppTestError(
            f"{method}() is not supported for {self.type}{key_part}. "
            "AppTest can inspect this element but does not implement "
            "interactions for it. Set its value through at.session_state if it "
            "has a key, or use a Playwright e2e test."
        )

    def run(self, *, timeout: float | None = None) -> AppTest:
        """Run the ``AppTest`` script which contains the element.

        Parameters
        ----------
        timeout
            The maximum number of seconds to run the script. None means
            use the AppTest's default.
        """
        return self.root.run(timeout=timeout)

    def __repr__(self) -> str:
        return util.repr_(self)


@dataclass(repr=False)
class UnknownElement(Element):
    def __init__(self, proto: ElementProto, root: ElementTree) -> None:
        ty = proto.WhichOneof("type")
        assert ty is not None
        self.proto = getattr(proto, ty)
        self.root = root
        self.type = ty
        proto_id = getattr(self.proto, "id", "") or ""
        self.key = user_key_from_element_id(proto_id) if proto_id else None

    @property
    def value(self) -> Any:
        proto_id = getattr(self.proto, "id", None)
        if proto_id:
            try:
                state = self.root.session_state
                if state is not None:
                    return state[proto_id]
            except (KeyError, ValueError):
                # Missing widget state or a non-widget id is expected for
                # unimplemented elements; fall back to a proto field.
                pass
        return _unknown_element_content(self.proto)


@dataclass(repr=False)
class Widget(Element, ABC):
    """Widget base class for testing."""

    id: str = field(repr=False)
    disabled: bool
    key: str | None
    _value: Any

    def __init__(self, proto: Any, root: ElementTree) -> None:
        self.proto = proto
        self.root = root
        self.key = user_key_from_element_id(self.id)
        self._value = None

    def _assert_can_interact(self) -> None:
        """Reject interactions that a browser user cannot perform."""
        if getattr(self.proto, "disabled", False):
            key_part = f" (key={self.key!r})" if self.key else ""
            raise AppTestError(
                f"Cannot update a disabled {self.type} widget{key_part}. "
                "A browser user cannot interact with a disabled widget."
            )

    def set_value(self, v: Any) -> Self:
        """Set the value of the widget."""
        self._assert_can_interact()
        self._value = v
        return self

    @property
    @abstractmethod
    def _widget_state(self) -> WidgetState: ...


El_co = TypeVar("El_co", bound=Element, covariant=True)


class ElementList(Generic[El_co]):
    def __init__(self, els: Sequence[El_co]) -> None:
        self._list: Sequence[El_co] = els

    def __len__(self) -> int:
        return len(self._list)

    @property
    def len(self) -> int:
        return len(self)

    @overload
    def __getitem__(self, idx: int) -> El_co: ...

    @overload
    def __getitem__(self, idx: slice) -> ElementList[El_co]: ...

    def __getitem__(self, idx: int | slice) -> El_co | ElementList[El_co]:
        if isinstance(idx, slice):
            return ElementList(self._list[idx])
        return self._list[idx]

    def __iter__(self) -> Iterator[El_co]:
        yield from self._list

    def __repr__(self) -> str:
        return util.repr_(self)

    def __eq__(self, other: ElementList[El_co] | object) -> bool:
        if isinstance(other, ElementList):
            return self._list == other._list
        return self._list == other

    def __hash__(self) -> int:
        return hash(tuple(self._list))

    def __call__(self, key: str) -> El_co:
        """Return the first element in this collection with the given user key.

        The same key can appear on different element types, so this returns the
        first match in this collection. Use ``get_by_key`` when an ambiguous
        key should raise.
        """
        for e in self._list:
            if getattr(e, "key", None) == key:
                return e
        raise KeyError(key)

    @property
    def values(self) -> Sequence[Any]:
        return [e.value for e in self]


W_co = TypeVar("W_co", bound=Widget, covariant=True)


class WidgetList(ElementList[W_co], Generic[W_co]):
    """ElementList narrowed to widgets for typing."""


class BlockList:
    """Sequence of layout blocks with optional lookup by user key."""

    def __init__(self, els: Sequence[Block]) -> None:
        self._list = list(els)

    def __len__(self) -> int:
        return len(self._list)

    @property
    def len(self) -> int:
        return len(self)

    @overload
    def __getitem__(self, idx: int) -> Block: ...

    @overload
    def __getitem__(self, idx: slice) -> BlockList: ...

    def __getitem__(self, idx: int | slice) -> Block | BlockList:
        if isinstance(idx, slice):
            return BlockList(self._list[idx])
        return self._list[idx]

    def __iter__(self) -> Iterator[Block]:
        return iter(self._list)

    def __repr__(self) -> str:
        return util.repr_(self)

    def __eq__(self, other: BlockList | object) -> bool:
        if isinstance(other, BlockList):
            return self._list == other._list
        return self._list == other

    def __hash__(self) -> int:
        return hash(tuple(self._list))

    def __call__(self, key: str) -> Block:
        """Return the first block in this collection with the given user key.

        The same key can appear on different block types, so this returns the
        first match in this collection. Use ``get_by_key`` when an ambiguous
        key should raise.
        """
        for e in self._list:
            if e.key == key:
                return e
        raise KeyError(key)


@dataclass(repr=False)
class AlertBase(Element):
    proto: AlertProto = field(repr=False)
    icon: str

    def __init__(self, proto: AlertProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root

    @property
    def value(self) -> str:
        return self.proto.body


@dataclass(repr=False)
class Error(AlertBase):
    def __init__(self, proto: AlertProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "error"


@dataclass(repr=False)
class Warning(AlertBase):  # noqa: A001
    def __init__(self, proto: AlertProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "warning"


@dataclass(repr=False)
class Info(AlertBase):
    def __init__(self, proto: AlertProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "info"


@dataclass(repr=False)
class Success(AlertBase):
    def __init__(self, proto: AlertProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "success"


@dataclass(repr=False)
class Button(Widget):
    """A representation of ``st.button`` and ``st.form_submit_button``."""

    _value: bool

    proto: ButtonProto = field(repr=False)
    label: str
    help: str
    form_id: str

    def __init__(self, proto: ButtonProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = False
        self.type = "button"

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        ws.trigger_value = self._value
        return ws

    @property
    def value(self) -> bool:
        """The value of the button. (bool)"""  # noqa: D400
        if self._value:
            return self._value
        state = self.root.session_state
        assert state
        return cast("bool", state[TESTING_KEY][self.id])

    def set_value(self, v: bool) -> Button:
        """Set the value of the button."""
        return super().set_value(v)

    def click(self) -> Button:
        """Set the value of the button to True."""
        return self.set_value(True)


@dataclass(repr=False)
class DownloadButton(Widget):
    """A representation of ``st.download_button``."""

    _value: bool

    proto: DownloadButtonProto = field(repr=False)
    label: str
    help: str
    form_id: str

    def __init__(self, proto: DownloadButtonProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = False
        self.type = "download_button"

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        ws.trigger_value = self._value
        return ws

    @property
    def value(self) -> bool:
        """The value of the download button. (bool)"""  # noqa: D400
        if self._value:
            return self._value
        state = self.root.session_state
        assert state
        return cast("bool", state[TESTING_KEY][self.id])

    def set_value(self, v: bool) -> DownloadButton:
        """Set the value of the download button."""
        return super().set_value(v)

    def click(self) -> DownloadButton:
        """Set the value of the download button to True."""
        return self.set_value(True)


@dataclass(repr=False)
class ChatInput(Widget):
    """A representation of ``st.chat_input``."""

    _value: str | None
    proto: ChatInputProto = field(repr=False)
    placeholder: str

    def __init__(self, proto: ChatInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "chat_input"

    def set_value(self, v: str | None) -> ChatInput:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        if self._value is not None:
            ws.chat_input_value.data = self._value
        return ws

    @property
    def value(self) -> str | ChatInputValue | None:
        """The pending or last submitted chat input value.

        Before ``.run()``, a pending ``set_value`` is that string (including ``""``).
        After ``.run()``, this is a plain ``str`` for a text-only chat input, a
        ``ChatInputValue`` when ``accept_file`` or ``accept_audio`` is enabled, and
        ``None`` when nothing was submitted.
        """
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        return state[TESTING_KEY][self.id]  # type: ignore


@dataclass(repr=False)
class Checkbox(Widget):
    """A representation of ``st.checkbox``."""

    _value: bool | None

    proto: CheckboxProto = field(repr=False)
    label: str
    help: str
    form_id: str

    def __init__(self, proto: CheckboxProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "checkbox"

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        ws.bool_value = self.value
        return ws

    @property
    def value(self) -> bool:
        """The value of the widget. (bool)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        return cast("bool", state[self.id])

    def set_value(self, v: bool) -> Checkbox:
        """Set the value of the widget."""
        return super().set_value(v)

    def check(self) -> Checkbox:
        """Set the value of the widget to True."""
        return self.set_value(True)

    def uncheck(self) -> Checkbox:
        """Set the value of the widget to False."""
        return self.set_value(False)


@dataclass(repr=False)
class Code(Element):
    """A representation of ``st.code``."""

    proto: CodeProto = field(repr=False)

    language: str
    show_line_numbers: bool
    key: None

    def __init__(self, proto: CodeProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "code"

    @property
    def value(self) -> str:
        """The value of the element. (str)"""  # noqa: D400
        return self.proto.code_text


@dataclass(repr=False)
class ColorPicker(Widget):
    """A representation of ``st.color_picker``."""

    _value: str | None
    label: str
    help: str
    form_id: str

    proto: ColorPickerProto = field(repr=False)

    def __init__(self, proto: ColorPickerProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "color_picker"

    @property
    def value(self) -> str:
        """The currently selected value as a hex string. (str)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        return cast("str", state[self.id])

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        ws.string_value = self.value
        return ws

    def set_value(self, v: str) -> ColorPicker:
        """Set the value of the widget as a hex string."""
        return super().set_value(v)

    def pick(self, v: str) -> ColorPicker:
        """Set the value of the widget as a hex string. May omit the "#" prefix."""
        if not v.startswith("#"):
            v = f"#{v}"
        return self.set_value(v)


@dataclass(repr=False)
class Dataframe(Element):
    proto: DataframeProto = field(repr=False)

    def __init__(self, proto: DataframeProto, root: ElementTree) -> None:
        self.proto = proto
        self.root = root
        self.type = "dataframe"
        # Extract user key from the element id if present
        self.key = user_key_from_element_id(proto.id) if proto.id else None

    @property
    def value(self) -> PandasDataframe:
        arrow_data = self.proto.arrow_data.data
        if (
            not arrow_data
            and self.proto.HasField("lazy_data")
            and self.proto.lazy_data.initial_chunk.data
        ):
            arrow_data = self.proto.lazy_data.initial_chunk.data
        return dataframe_util.convert_arrow_bytes_to_pandas_df(arrow_data)


SingleDateValue: TypeAlias = date | datetime
DateValue: TypeAlias = SingleDateValue | Sequence[SingleDateValue] | None


@dataclass(repr=False)
class DateInput(Widget):
    """A representation of ``st.date_input``."""

    _value: DateValue | InitialValue | None
    proto: DateInputProto = field(repr=False)
    label: str
    min: date
    max: date
    is_range: bool
    help: str
    form_id: str

    def __init__(self, proto: DateInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "date_input"
        self.min = datetime.strptime(proto.min, "%Y-%m-%d").date()
        self.max = datetime.strptime(proto.max, "%Y-%m-%d").date()

    def set_value(self, v: DateValue) -> DateInput:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id

        serde = DateInputSerde(None)  # type: ignore
        ws.string_array_value.data[:] = serde.serialize(self.value)
        return ws

    @property
    def value(self) -> DateWidgetReturn:
        """The value of the widget. (date or Tuple of date)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            parsed, _ = _parse_date_value(self._value)
            return tuple(parsed) if parsed is not None else None  # type: ignore
        state = self.root.session_state
        assert state
        return state[self.id]  # type: ignore


@dataclass(repr=False)
class Exception(Element):  # noqa: A001
    message: str
    is_markdown: bool
    stack_trace: list[str]
    is_warning: bool

    def __init__(self, proto: ExceptionProto, root: ElementTree) -> None:
        self.key = None
        self.root = root
        self.proto = proto
        self.type = "exception"

        self.is_markdown = proto.message_is_markdown
        self.stack_trace = list(proto.stack_trace)

    @property
    def value(self) -> str:
        return self.message


@dataclass(repr=False)
class HeadingBase(Element, ABC):
    proto: HeadingProto = field(repr=False)

    tag: str
    anchor: str | None
    hide_anchor: bool
    key: None

    def __init__(self, proto: HeadingProto, root: ElementTree, type_: str) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = type_

    @property
    def value(self) -> str:
        return self.proto.body

    @property
    def icon(self) -> str:
        """The heading icon, or empty string when none is set."""
        return self.proto.icon


@dataclass(repr=False)
class Header(HeadingBase):
    def __init__(self, proto: HeadingProto, root: ElementTree) -> None:
        super().__init__(proto, root, "header")


@dataclass(repr=False)
class Subheader(HeadingBase):
    def __init__(self, proto: HeadingProto, root: ElementTree) -> None:
        super().__init__(proto, root, "subheader")


@dataclass(repr=False)
class Title(HeadingBase):
    def __init__(self, proto: HeadingProto, root: ElementTree) -> None:
        super().__init__(proto, root, "title")


@dataclass(repr=False)
class Image(Element):
    """A representation of ``st.image``."""

    proto: ImageListProto = field(repr=False)
    key: None

    def __init__(self, proto: ImageListProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "image"

    @property
    def value(self) -> list[str]:
        """The image URLs for this element."""
        return [img.url for img in self.proto.imgs]

    @property
    def captions(self) -> list[str]:
        """The image captions for this element."""
        return [img.caption for img in self.proto.imgs]


@dataclass(repr=False)
class Json(Element):
    proto: JsonProto = field(repr=False)

    expanded: bool

    def __init__(self, proto: JsonProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "json"

    @property
    def value(self) -> str:
        return self.proto.body


@dataclass(repr=False)
class Markdown(Element):
    proto: MarkdownProto = field(repr=False)

    allow_html: bool
    key: None

    def __init__(self, proto: MarkdownProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "markdown"

    @property
    def value(self) -> str:
        return self.proto.body

    @property
    def is_caption(self) -> bool:
        """Whether this is a caption element (derived from element_type)."""
        return self.proto.element_type == MarkdownProto.Type.CAPTION


@dataclass(repr=False)
class Caption(Markdown):
    def __init__(self, proto: MarkdownProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "caption"


@dataclass(repr=False)
class Divider(Markdown):
    def __init__(self, proto: MarkdownProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "divider"


@dataclass(repr=False)
class Latex(Markdown):
    def __init__(self, proto: MarkdownProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "latex"


@dataclass(repr=False)
class Space(Element):
    """A representation of ``st.space``."""

    proto: SpaceProto = field(repr=False)
    key: None
    size: SpaceSize | None

    def __init__(
        self,
        proto: SpaceProto,
        root: ElementTree,
        *,
        size: SpaceSize | None,
    ) -> None:
        super().__init__(proto, root)
        self.key = None
        self.type = "space"
        self.size = size

    @property
    def value(self) -> SpaceSize | None:
        """The ``size`` argument passed to ``st.space`` (``"small"`` when omitted).

        This is ``None`` only if the size could not be reconstructed from the
        element's width config.
        """
        return self.size


@dataclass(repr=False)
class Metric(Element):
    proto: MetricProto
    label: str
    delta: str
    color: str
    help: str

    def __init__(self, proto: MetricProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "metric"

    @property
    def value(self) -> str:
        return self.proto.body


@dataclass(repr=False)
class ButtonGroup(Widget, Generic[T]):
    """A representation of ``st.pills`` and ``st.segmented_control``."""

    _value: T | list[T] | InitialValue | None

    proto: ButtonGroupProto = field(repr=False)
    options: list[str]
    form_id: str

    def __init__(self, proto: ButtonGroupProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "button_group"
        # Store formatted content strings for value serialization
        self.options = [opt.content for opt in proto.options]

    @property
    def _is_single_select(self) -> bool:
        """Check if this is a single-select widget."""
        return self.proto.click_mode == ButtonGroupProto.ClickMode.SINGLE_SELECT

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        ws.string_array_value.data[:] = self.formatted_values
        return ws

    @property
    def value(self) -> T | list[T] | None:
        """The currently selected value(s) from the options.

        For single-select mode, returns a single value (or None if nothing selected).
        For multi-select mode, returns a list of values.
        """
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        return cast("T | list[T]", state[self.id])

    @property
    def indices(self) -> Sequence[int]:
        """The indices of the currently selected values from the options. (list)"""  # noqa: D400
        return [self.options.index(v) for v in self.formatted_values]

    @property
    def formatted_values(self) -> Sequence[str]:
        """The formatted string values for the current selection."""
        format_func = self.format_func
        value = self.value
        if value is None:
            return []
        if self._is_single_select:
            return [_format_value_for_widget(format_func, value)]
        return [
            _format_value_for_widget(format_func, v) for v in cast("list[T]", value)
        ]

    @property
    def format_func(self) -> Callable[[Any], Any]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        return cast("Callable[[Any], Any]", ss[TESTING_KEY][self.id])

    def set_value(self, v: T | list[T] | None) -> ButtonGroup[T]:
        """Set the value of the widget.

        For single-select mode, pass a single value or None to clear the selection.
        For multi-select mode, pass a list of values (use empty list to clear).
        """
        return super().set_value(v)

    def select(self, v: T) -> ButtonGroup[T]:
        """Add a selection to the widget.

        For single-select mode, this sets the value to v.
        For multi-select mode, this adds v to the selection if not already selected.
        """
        if self._is_single_select:
            return self.set_value(v)
        current = list(cast("list[T]", self.value) or [])
        if v in current:
            return self
        new = current.copy()
        new.append(v)
        return self.set_value(new)

    def unselect(self, v: T) -> ButtonGroup[T]:
        """Remove a selection from the widget.

        For single-select mode, this sets the value to None if v is the current value.
        For multi-select mode, this removes v from the selection.
        """
        if self._is_single_select:
            if self.value == v:
                return self.set_value(None)
            return self
        current = list(cast("list[T]", self.value) or [])
        if v not in current:
            return self
        new = current.copy()
        while v in new:
            new.remove(v)
        return self.set_value(new)


@dataclass(repr=False)
class Feedback(Widget):
    """A representation of ``st.feedback``."""

    _value: int | InitialValue | None

    proto: FeedbackProto = field(repr=False)
    form_id: str

    def __init__(self, proto: FeedbackProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "feedback"

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.

        Uses string_value as wire format to distinguish three states:
        - None: User explicitly cleared -> string_value = ""
        - int: User selected -> string_value = str(value)
        - No string_value set: No interaction yet (use default)
        """
        ws = WidgetState()
        ws.id = self.id

        # Get the effective value: either from explicit set_value() or session state
        effective_value = self.value

        if effective_value is None:
            ws.string_value = ""  # Cleared or no default
        else:
            ws.string_value = str(effective_value)  # User selected a value
        return ws

    @property
    def value(self) -> int | None:
        """The currently selected feedback value. (int or None)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        return cast("int | None", state[self.id])

    def set_value(self, v: int | None) -> Feedback:
        """Set the value of the feedback widget. (int or None)"""  # noqa: D400
        return super().set_value(v)


@dataclass(repr=False)
class FileUploader(Widget):
    r"""A representation of ``st.file_uploader``.

    Files are provided as tuples of (filename, content, mime_type).

    Example
    -------
    >>> at = AppTest.from_string('''
    ...     import streamlit as st
    ...     uploaded = st.file_uploader("Upload a file")
    ...     if uploaded:
    ...         st.write(f"Uploaded: {uploaded.name}")
    ... ''')
    >>> at.run()
    >>> at.file_uploader[0].set_value(("data.csv", b"col1,col2\n1,2", "text/csv"))
    >>> at.run()
    >>> at.markdown[0].value
    'Uploaded: data.csv'
    """

    # Stores list of (file_id, filename, content, mime_type) tuples
    # InitialValue means no explicit set_value/upload/clear was called
    _files: list[tuple[str, str, bytes, str]] | InitialValue | None

    proto: FileUploaderProto = field(repr=False)
    label: str
    help: str
    form_id: str

    def __init__(self, proto: FileUploaderProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._files = InitialValue()
        self.type = "file_uploader"

    @property
    def accept_multiple_files(self) -> bool:
        """Whether multiple files can be uploaded. (bool)"""  # noqa: D400
        return self.proto.multiple_files

    @property
    def accept_directory(self) -> bool:
        """Whether directory uploads are accepted. (bool)"""  # noqa: D400
        return self.proto.accept_directory

    @property
    def allowed_type(self) -> list[str]:
        """Allowed file types for upload. (list of str)"""  # noqa: D400
        return list(self.proto.type)

    def set_value(  # ty: ignore[invalid-method-override]
        self,
        files: (tuple[str, bytes, str] | Sequence[tuple[str, bytes, str]] | None),
    ) -> Self:
        """Set the uploaded file(s) for testing.

        Parameters
        ----------
        files
            A tuple of (filename, content, mime_type) for single file upload,
            or a sequence of such tuples for multiple file upload.
            Set to ``None`` to clear uploaded files.

        Returns
        -------
        FileUploader
            The FileUploader instance for method chaining.
        """
        from uuid import uuid4

        self._assert_can_interact()
        if files is None:
            self._files = None
        elif isinstance(files, tuple) and len(files) == 3 and isinstance(files[0], str):
            # Single file as tuple (filename, content, mime_type)
            single_file = cast(  # type: ignore[redundant-cast]
                "tuple[str, bytes, str]", files
            )
            filename, content, mime_type = single_file
            self._files = [(str(uuid4()), filename, content, mime_type)]
        else:
            # Multiple files
            files_list = cast("Sequence[tuple[str, bytes, str]]", files)
            self._files = [
                (str(uuid4()), filename, content, mime_type)
                for filename, content, mime_type in files_list
            ]
        return self

    def upload(
        self,
        filename: str,
        content: bytes,
        mime_type: str = "application/octet-stream",
    ) -> Self:
        """Upload a single file for testing.

        Parameters
        ----------
        filename
            The name of the file.
        content
            The file content as bytes.
        mime_type
            The MIME type of the file. Defaults to "application/octet-stream".

        Returns
        -------
        FileUploader
            The FileUploader instance for method chaining.
        """
        from uuid import uuid4

        self._assert_can_interact()
        if self._files is None or isinstance(self._files, InitialValue):
            self._files = []
        self._files.append((str(uuid4()), filename, content, mime_type))
        return self

    def clear(self) -> Self:
        """Clear all uploaded files.

        Returns
        -------
        FileUploader
            The FileUploader instance for method chaining.
        """
        self._assert_can_interact()
        self._files = None
        return self

    def _get_files_to_register(self) -> list[tuple[str, str, bytes, str]]:
        """Return files to register: list of (file_id, filename, content, mime_type).

        If no explicit set_value/upload/clear was called, derive from existing
        UploadedFile(s) in session_state to persist files across runs.
        """
        # If explicitly set, use that value
        if not isinstance(self._files, InitialValue):
            return self._files or []

        # Fall back to existing UploadedFile(s) in session_state
        from streamlit.runtime.uploaded_file_manager import UploadedFile

        state = self.root.session_state
        if not state:
            return []

        try:
            current_value = state[self.id]
        except KeyError:
            return []

        if current_value is None:
            return []

        # Handle both single file and multiple files
        files_list: list[UploadedFile] = []
        if isinstance(current_value, list):
            files_list = current_value
        elif isinstance(current_value, UploadedFile):
            files_list = [current_value]

        return [
            (f.file_id, f.name, f.getvalue(), f.type)
            for f in files_list
            if isinstance(f, UploadedFile)
        ]

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget."""
        from streamlit.proto.Common_pb2 import (
            FileUploaderState as FileUploaderStateProto,
        )

        ws = WidgetState()
        ws.id = self.id

        # Use _get_files_to_register which handles fallback to session_state
        files_to_use = self._get_files_to_register()

        if not files_to_use:
            # Return empty state only if explicitly cleared (not InitialValue)
            # or if there are no files in session_state
            return ws

        # Create file uploader state proto with pre-generated file IDs
        state_proto = FileUploaderStateProto()

        for file_id, filename, content, _mime_type in files_to_use:
            file_info = state_proto.uploaded_file_info.add()
            file_info.file_id = file_id
            file_info.name = filename
            file_info.size = len(content)
            file_info.file_urls.file_id = file_id
            file_info.file_urls.upload_url = f"/mock/upload/test session id/{file_id}"
            file_info.file_urls.delete_url = f"/mock/upload/test session id/{file_id}"

        ws.file_uploader_state_value.CopyFrom(state_proto)
        return ws

    @property
    def value(self) -> Any:
        """The current uploaded file(s).

        Returns the UploadedFile object(s) or None, depending on the
        ``accept_multiple_files`` setting.
        """
        state = self.root.session_state
        assert state
        return state[self.id]


@dataclass(repr=False)
class MenuButton(Widget, Generic[T]):
    """A representation of ``st.menu_button``."""

    _value: T | None

    proto: MenuButtonProto = field(repr=False)
    label: str
    options: list[str]
    help: str

    def __init__(self, proto: MenuButtonProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = None
        self.type = "menu_button"
        self.options = list(proto.options)

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        if self._value is not None:
            try:
                ws.string_trigger_value.data = self.format_func(self._value)
            except builtins.Exception:
                ws.string_trigger_value.data = str(self._value)
        return ws

    @property
    def value(self) -> T | None:
        """The selected option value, or None if no option was clicked. (Any)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        # For trigger widgets, the value is stored in TESTING_KEY along with format_func
        testing_data = state[TESTING_KEY][self.id]
        return cast("T | None", testing_data["value"])

    @property
    def format_func(self) -> Callable[[Any], str]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        testing_data = ss[TESTING_KEY][self.id]
        return cast("Callable[[Any], str]", testing_data["format_func"])

    def set_value(self, v: T | None) -> MenuButton[T]:
        """Set the selected option value."""
        return super().set_value(v)

    def click(self, v: T) -> MenuButton[T]:
        """Click an option by value, simulating user selection."""
        return self.set_value(v)

    def click_index(self, index: int) -> MenuButton[T]:
        """Click an option by index, simulating user selection."""
        # Use original unformatted options from testing data to preserve the correct type
        ss = self.root.session_state
        testing_data = ss[TESTING_KEY][self.id]
        original_options = testing_data["options"]
        return self.set_value(cast("T", original_options[index]))


@dataclass(repr=False)
class Multiselect(Widget, Generic[T]):
    """A representation of ``st.multiselect``."""

    _value: list[T] | None

    proto: MultiSelectProto = field(repr=False)
    label: str
    options: list[str]
    max_selections: int
    help: str
    form_id: str

    def __init__(self, proto: MultiSelectProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "multiselect"
        self.options = list(proto.options)

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        ws.string_array_value.data[:] = self.values
        return ws

    @property
    def value(self) -> list[T]:
        """The currently selected values from the options. (list)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        return cast("list[T]", state[self.id])

    @property
    def indices(self) -> Sequence[int]:
        """The indices of the currently selected values from the options. (list)"""  # noqa: D400
        return [self.options.index(v) for v in self.values]

    @property
    def values(self) -> Sequence[str]:
        """The currently selected values from the options. (list)"""  # noqa: D400
        format_func = self.format_func
        return [_format_value_for_widget(format_func, v) for v in self.value]

    @property
    def format_func(self) -> Callable[[Any], Any]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        return cast("Callable[[Any], Any]", ss[TESTING_KEY][self.id])

    def set_value(self, v: list[T]) -> Multiselect[T]:
        """Set the value of the multiselect widget. (list)"""  # noqa: D400
        return super().set_value(v)

    def select(self, v: T) -> Multiselect[T]:
        """
        Add a selection to the widget. Do nothing if the value is already selected.\
        If testing a multiselect widget with repeated options, use ``set_value``\
        instead.
        """
        current = self.value
        if v in current:
            return self
        new = current.copy()
        new.append(v)
        self.set_value(new)
        return self

    def unselect(self, v: T) -> Multiselect[T]:
        """
        Remove a selection from the widget. Do nothing if the value is not\
        already selected. If a value is selected multiple times, the first\
        instance is removed.
        """
        current = self.value
        if v not in current:
            return self
        new = current.copy()
        while v in new:
            new.remove(v)
        self.set_value(new)
        return self


Number: TypeAlias = int | float


@dataclass(repr=False)
class NumberInput(Widget):
    """A representation of ``st.number_input``."""

    _value: Number | InitialValue | None
    proto: NumberInputProto = field(repr=False)
    label: str
    min: Number | None
    max: Number | None
    step: Number
    help: str
    form_id: str

    def __init__(self, proto: NumberInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "number_input"
        self.min = proto.min if proto.has_min else None
        self.max = proto.max if proto.has_max else None

    def set_value(self, v: Number | None) -> NumberInput:
        """Set the value of the ``st.number_input`` widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        if self.value is not None:
            ws.double_value = self.value
        return ws

    @property
    def value(self) -> Number | None:
        """The current value of the ``st.number_input`` widget."""
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state

        # Awkward to do this with `cast`
        return state[self.id]  # type: ignore

    def increment(self) -> NumberInput:
        """Increment the ``st.number_input`` widget as if the user clicked "+"."""
        if self.value is None:
            return self

        v = min(self.value + self.step, self.max or float("inf"))
        return self.set_value(v)

    def decrement(self) -> NumberInput:
        """Decrement the ``st.number_input`` widget as if the user clicked "-"."""
        if self.value is None:
            return self

        v = max(self.value - self.step, self.min or float("-inf"))
        return self.set_value(v)


@dataclass(repr=False)
class Radio(Widget, Generic[T]):
    """A representation of ``st.radio``."""

    _value: T | InitialValue | None

    proto: RadioProto = field(repr=False)
    label: str
    options: list[str]
    horizontal: bool
    help: str
    form_id: str

    def __init__(self, proto: RadioProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "radio"
        self.options = list(proto.options)

    @property
    def index(self) -> int | None:
        """The index of the current selection. (int)"""  # noqa: D400
        if self.value is None:
            return None
        formatted_value = _format_value_for_widget(self.format_func, self.value)
        return self.options.index(formatted_value)

    @property
    def value(self) -> T | None:
        """The currently selected value from the options. (Any)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        return cast("T", state[self.id])

    @property
    def format_func(self) -> Callable[[Any], Any]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        return cast("Callable[[Any], Any]", ss[TESTING_KEY][self.id])

    def set_value(self, v: T | None) -> Radio[T]:
        """Set the selection by value."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        if self.index is not None and len(self.options) > 0:
            ws.string_value = self.options[self.index]
        return ws


@dataclass(repr=False)
class Selectbox(Widget, Generic[T]):
    """A representation of ``st.selectbox``."""

    _value: T | InitialValue | None

    proto: SelectboxProto = field(repr=False)
    label: str
    options: list[str]
    help: str
    form_id: str

    def __init__(self, proto: SelectboxProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "selectbox"
        self.options = list(proto.options)

    @property
    def index(self) -> int | None:
        """The index of the current selection. (int)"""  # noqa: D400
        if self.value is None:
            return None

        if len(self.options) == 0:
            return 0
        formatted_value = _format_value_for_widget(self.format_func, self.value)
        return self.options.index(formatted_value)

    @property
    def value(self) -> T | None:
        """The currently selected value from the options. (Any)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        return cast("T", state[self.id])

    @property
    def format_func(self) -> Callable[[Any], Any]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        return cast("Callable[[Any], Any]", ss[TESTING_KEY][self.id])

    def set_value(self, v: T | None) -> Selectbox[T]:
        """Set the selection by value."""
        return super().set_value(v)

    def select(self, v: T | None) -> Selectbox[T]:
        """Set the selection by value."""
        return self.set_value(v)

    def select_index(self, index: int | None) -> Selectbox[T]:
        """Set the selection by index."""
        if index is None:
            return self.set_value(None)
        return self.set_value(cast("T", self.options[index]))

    @property
    def _widget_state(self) -> WidgetState:
        """Protobuf message representing the state of the widget, including
        any interactions that have happened.
        Should be the same as the frontend would produce for those interactions.
        """
        ws = WidgetState()
        ws.id = self.id
        if self.index is not None and len(self.options) > 0:
            ws.string_value = self.options[self.index]
        return ws


@dataclass(repr=False)
class SelectSlider(Widget, Generic[T]):
    """A representation of ``st.select_slider``."""

    _value: T | Sequence[T] | None

    proto: SliderProto = field(repr=False)
    label: str
    data_type: SliderProto.DataType.ValueType
    options: list[str]
    help: str
    form_id: str

    def __init__(self, proto: SliderProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "select_slider"
        self.options = list(proto.options)

    def set_value(self, v: T | Sequence[T]) -> SelectSlider[T]:
        """Set the (single) selection by value."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        # Build formatted options mapping
        # Note: self.options already contains formatted strings from the proto,
        # so we should NOT apply format_func to them again
        format_func = self.format_func
        formatted_option_to_index = {opt: idx for idx, opt in enumerate(self.options)}

        # Determine if this is a range value
        is_range = isinstance(self.value, (list, tuple)) and len(self.value) == 2

        serde = SelectSliderSerde(
            self.options,
            formatted_option_to_index=formatted_option_to_index,
            default_indices=[0] if not is_range else [0, len(self.options) - 1],
            format_func=format_func,
        )

        try:
            if is_range:
                v = serde.serialize(tuple(self.value))  # type: ignore
            else:
                v = serde.serialize(self.value)  # type: ignore
        except (ValueError, TypeError) as e:
            raise ValueError(f"Could not serialize value {self.value}") from e

        ws = WidgetState()
        ws.id = self.id
        ws.string_array_value.data[:] = v
        return ws

    @property
    def value(self) -> T | Sequence[T]:
        """The currently selected value or range. (Any or Sequence of Any)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        # Awkward to do this with `cast`
        return state[self.id]  # type: ignore

    @property
    def format_func(self) -> Callable[[Any], Any]:
        """The widget's formatting function for displaying options. (callable)"""  # noqa: D400
        ss = self.root.session_state
        return cast("Callable[[Any], Any]", ss[TESTING_KEY][self.id])

    def set_range(self, lower: T, upper: T) -> SelectSlider[T]:
        """Set the ranged selection by values."""
        return self.set_value([lower, upper])


SliderValueT = TypeVar("SliderValueT", int, float, date, time, datetime)


@dataclass(repr=False)
class Slider(Widget, Generic[SliderValueT]):
    """A representation of ``st.slider``."""

    _value: SliderValueT | Sequence[SliderValueT] | None

    proto: SliderProto = field(repr=False)
    label: str
    data_type: SliderProto.DataType.ValueType
    min: SliderValueT
    max: SliderValueT
    step: SliderStep
    help: str
    form_id: str

    def __init__(self, proto: SliderProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self.type = "slider"

    def set_value(
        self, v: SliderValueT | Sequence[SliderValueT]
    ) -> Slider[SliderValueT]:
        """Set the (single) value of the slider."""
        return cast("Slider[SliderValueT]", super().set_value(v))

    @property
    def _widget_state(self) -> WidgetState:
        data_type = self.proto.data_type
        serde = SliderSerde([], data_type, True, None, self.proto.min, self.proto.max)
        v = serde.serialize(self.value)

        ws = WidgetState()
        ws.id = self.id
        ws.double_array_value.data[:] = v
        return ws

    @property
    def value(self) -> SliderValueT | Sequence[SliderValueT]:
        """The currently selected value or range. (Any or Sequence of Any)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        # Awkward to do this with `cast`
        return state[self.id]  # type: ignore

    def set_range(
        self, lower: SliderValueT, upper: SliderValueT
    ) -> Slider[SliderValueT]:
        """Set the ranged value of the slider."""
        return self.set_value([lower, upper])


@dataclass(repr=False)
class Table(Element):
    proto: TableProto = field(repr=False)

    def __init__(self, proto: TableProto, root: ElementTree) -> None:
        self.key = None
        self.proto = proto
        self.root = root
        self.type = "table"

    @property
    def value(self) -> PandasDataframe:
        return dataframe_util.convert_arrow_bytes_to_pandas_df(
            self.proto.arrow_data.data
        )


@dataclass(repr=False)
class Text(Element):
    proto: TextProto = field(repr=False)

    key: None = None

    def __init__(self, proto: TextProto, root: ElementTree) -> None:
        self.proto = proto
        self.root = root
        self.type = "text"

    @property
    def value(self) -> str:
        """The value of the element. (str)"""  # noqa: D400
        return self.proto.body


@dataclass(repr=False)
class TextArea(Widget):
    """A representation of ``st.text_area``."""

    _value: str | InitialValue | None

    proto: TextAreaProto = field(repr=False)
    label: str
    max_chars: int
    placeholder: str
    help: str
    form_id: str

    def __init__(self, proto: TextAreaProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "text_area"

    def set_value(self, v: str | None) -> TextArea:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        if self.value is not None:
            ws.string_value = self.value
        return ws

    @property
    def value(self) -> str | None:
        """The current value of the widget. (str)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        # Awkward to do this with `cast`
        return state[self.id]  # type: ignore

    def input(self, v: str) -> TextArea:
        """
        Set the value of the widget only if the value does not exceed the\
        maximum allowed characters.
        """
        # TODO: should input be setting or appending?
        if self.max_chars and len(v) > self.max_chars:
            return self
        return self.set_value(v)


@dataclass(repr=False)
class TextInput(Widget):
    """A representation of ``st.text_input``."""

    _value: str | InitialValue | None
    proto: TextInputProto = field(repr=False)
    label: str
    max_chars: int
    autocomplete: str
    placeholder: str
    help: str
    form_id: str

    def __init__(self, proto: TextInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "text_input"

    def set_value(self, v: str | None) -> TextInput:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        if self.value is not None:
            ws.string_value = self.value
        return ws

    @property
    def value(self) -> str | None:
        """The current value of the widget. (str)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        # Awkward to do this with `cast`
        return state[self.id]  # type: ignore

    def input(self, v: str) -> TextInput:
        """
        Set the value of the widget only if the value does not exceed the\
        maximum allowed characters.
        """
        # TODO: should input be setting or appending?
        if self.max_chars and len(v) > self.max_chars:
            return self
        return self.set_value(v)


TimeValue: TypeAlias = time | datetime
DateTimeWidgetValue: TypeAlias = datetime


@dataclass(repr=False)
class TimeInput(Widget):
    """A representation of ``st.time_input``."""

    _value: TimeValue | InitialValue | None
    proto: TimeInputProto = field(repr=False)
    label: str
    step: int
    help: str
    form_id: str

    def __init__(self, proto: TimeInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "time_input"

    def set_value(self, v: TimeValue | None) -> TimeInput:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id

        serde = TimeInputSerde(None, step=self.step)
        serialized_value = serde.serialize(self.value)
        if serialized_value is not None:
            ws.string_value = serialized_value
        return ws

    @property
    def value(self) -> time | None:
        """The current value of the widget. (time)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            v = self._value
            return v.time() if isinstance(v, datetime) else v
        state = self.root.session_state
        assert state
        return state[self.id]  # type: ignore

    def increment(self) -> TimeInput:
        """Select the next available time."""
        if self.value is None:
            return self
        dt = datetime.combine(date.today(), self.value) + timedelta(seconds=self.step)
        return self.set_value(dt.time())

    def decrement(self) -> TimeInput:
        """Select the previous available time."""
        if self.value is None:
            return self
        dt = datetime.combine(date.today(), self.value) - timedelta(seconds=self.step)
        return self.set_value(dt.time())


@dataclass(repr=False)
class DateTimeInput(Widget):
    """A representation of ``st.datetime_input``."""

    _value: DateTimeWidgetValue | InitialValue | None
    proto: DateTimeInputProto = field(repr=False)
    label: str
    format: str
    min: str
    max: str
    step: int
    help: str
    form_id: str

    def __init__(self, proto: DateTimeInputProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = InitialValue()
        self.type = "date_time_input"

    def set_value(self, v: DateTimeWidgetValue | None) -> DateTimeInput:
        """Set the value of the widget."""
        return super().set_value(v)

    @property
    def _widget_state(self) -> WidgetState:
        from datetime import datetime

        def _parse_dt(value: str) -> datetime:
            try:
                return datetime.strptime(value, "%Y-%m-%dT%H:%M")
            except ValueError:
                return datetime.strptime(value, "%Y/%m/%d, %H:%M")

        ws = WidgetState()
        ws.id = self.id

        # Parse min and max values for validation
        min_dt = _parse_dt(self.min)
        max_dt = _parse_dt(self.max)

        serde = DateTimeInputSerde(value=None, min=min_dt, max=max_dt)
        serialized_value = serde.serialize(self.value)
        if serialized_value is not None:
            ws.string_array_value.data[:] = serialized_value
        return ws

    @property
    def value(self) -> datetime | None:
        """The current value of the widget. (datetime)"""  # noqa: D400
        if not isinstance(self._value, InitialValue):
            return self._value
        state = self.root.session_state
        assert state
        return state[self.id]  # type: ignore


@dataclass(repr=False)
class Toast(Element):
    proto: ToastProto = field(repr=False)
    icon: str

    def __init__(self, proto: ToastProto, root: ElementTree) -> None:
        self.proto = proto
        self.key = None
        self.root = root
        self.type = "toast"

    @property
    def value(self) -> str:
        return self.proto.body


@dataclass(repr=False)
class Toggle(Widget):
    """A representation of ``st.toggle``."""

    _value: bool | None

    proto: CheckboxProto = field(repr=False)
    label: str
    help: str
    form_id: str

    def __init__(self, proto: CheckboxProto, root: ElementTree) -> None:
        super().__init__(proto, root)
        self._value = None
        self.type = "toggle"

    @property
    def _widget_state(self) -> WidgetState:
        ws = WidgetState()
        ws.id = self.id
        ws.bool_value = self.value
        return ws

    @property
    def value(self) -> bool:
        """The current value of the widget. (bool)"""  # noqa: D400
        if self._value is not None:
            return self._value
        state = self.root.session_state
        assert state
        return cast("bool", state[self.id])

    def set_value(self, v: bool) -> Toggle:
        """Set the value of the widget."""
        return super().set_value(v)


@dataclass(repr=False)
class Block:
    """A container of other elements.

    Elements within a Block can be inspected and interacted with. This follows
    the same syntax as inspecting and interacting within an ``AppTest`` object.

    For all container classes, parameters of the original element can be
    obtained as properties. For example, ``ChatMessage.avatar`` and
    ``Tab.label``.
    """

    type: str
    children: dict[int, Node]
    proto: Any = field(repr=False)
    root: ElementTree = field(repr=False)
    _block_id: str = field(repr=False, default="", init=False)

    def __init__(
        self,
        proto: BlockProto | None,
        root: ElementTree,
    ) -> None:
        self.children = {}
        self.proto = proto
        if proto:
            ty = proto.WhichOneof("type")
            if ty is not None:
                self.type = ty
            else:
                # `st.container` has no sub-message
                self.type = "container"
        else:
            self.type = "unknown"
        self.root = root
        self._block_id = getattr(proto, "id", "") or "" if proto is not None else ""

    def __len__(self) -> int:
        return len(self.children)

    def __iter__(self) -> Iterator[Node]:
        yield self
        for child_idx in self.children:
            yield from self.children[child_idx]

    def __getitem__(self, k: int) -> Node:
        return self.children[k]

    @property
    def key(self) -> str | None:
        """User key for this block, if the corresponding command set one."""
        proto = self.proto
        block_id = self._block_id or (
            getattr(proto, "id", None) if proto is not None else None
        )
        if block_id:
            return user_key_from_element_id(block_id)
        if self.type == "form" and proto is not None:
            return cast("str", proto.form.form_id)
        return None

    def get_by_key(self, key: str) -> Node:
        """Return the unique current node with this user key.

        Works for widgets, keyed display elements, and keyed containers.
        Raises ``KeyError`` if no node matches and ``AppTestError`` if
        more than one node matches.
        """
        matches = [e for e in self if e is not self and getattr(e, "key", None) == key]
        if len(matches) == 1:
            return matches[0]
        if not matches:
            raise KeyError(key)
        raise AppTestError(
            f"Multiple elements have key {key!r}. "
            "Scope the query to a container or use a typed collection."
        )

    # We could implement these using __getattr__ but that would have
    # much worse type information.
    @property
    def button(self) -> WidgetList[Button]:
        return WidgetList(self.get("button"))  # type: ignore

    @property
    def button_group(self) -> WidgetList[ButtonGroup[Any]]:
        return WidgetList(self.get("button_group"))  # type: ignore

    @property
    def pills(self) -> WidgetList[ButtonGroup[Any]]:
        """st.pills widgets (subset of button_group with PILLS style)."""
        return WidgetList(
            [
                button_group
                for button_group in self.button_group
                if button_group.proto.style == ButtonGroupProto.Style.PILLS
            ]
        )

    @property
    def segmented_control(self) -> WidgetList[ButtonGroup[Any]]:
        """st.segmented_control widgets (subset of button_group with SEGMENTED_CONTROL style)."""
        return WidgetList(
            [
                button_group
                for button_group in self.button_group
                if button_group.proto.style == ButtonGroupProto.Style.SEGMENTED_CONTROL
            ]
        )

    @property
    def caption(self) -> ElementList[Caption]:
        return ElementList(self.get("caption"))  # type: ignore

    @property
    def chat_input(self) -> WidgetList[ChatInput]:
        return WidgetList(self.get("chat_input"))  # type: ignore

    @property
    def chat_message(self) -> Sequence[ChatMessage]:
        return self.get("chat_message")  # type: ignore

    @property
    def checkbox(self) -> WidgetList[Checkbox]:
        return WidgetList(self.get("checkbox"))  # type: ignore

    @property
    def code(self) -> ElementList[Code]:
        return ElementList(self.get("code"))  # type: ignore

    @property
    def color_picker(self) -> WidgetList[ColorPicker]:
        return WidgetList(self.get("color_picker"))  # type: ignore

    @property
    def columns(self) -> Sequence[Column]:
        return self.get("column")  # type: ignore

    @property
    def container(self) -> BlockList:
        """``st.container`` blocks, including horizontal/flex containers.

        The implicit row wrapper created by ``st.columns`` is excluded.
        """
        return BlockList(
            [
                e
                for e in self
                if isinstance(e, Block)
                and e is not self
                and e.type in {"container", "flex_container"}
                # st.columns currently emits a flex-container whose direct
                # children are Column blocks; that wrapper is not st.container.
                and not any(isinstance(child, Column) for child in e.children.values())
            ]
        )

    @property
    def dataframe(self) -> ElementList[Dataframe]:
        return ElementList(self.get("dataframe"))  # type: ignore

    @property
    def date_input(self) -> WidgetList[DateInput]:
        return WidgetList(self.get("date_input"))  # type: ignore

    @property
    def datetime_input(self) -> WidgetList[DateTimeInput]:
        return WidgetList(self.get("date_time_input"))  # type: ignore

    @property
    def divider(self) -> ElementList[Divider]:
        return ElementList(self.get("divider"))  # type: ignore

    @property
    def download_button(self) -> WidgetList[DownloadButton]:
        return WidgetList(self.get("download_button"))  # type: ignore

    @property
    def error(self) -> ElementList[Error]:
        return ElementList(self.get("error"))  # type: ignore

    @property
    def exception(self) -> ElementList[Exception]:
        return ElementList(self.get("exception"))  # type: ignore

    @property
    def feedback(self) -> WidgetList[Feedback]:
        return WidgetList(self.get("feedback"))  # type: ignore

    @property
    def file_uploader(self) -> WidgetList[FileUploader]:
        return WidgetList(self.get("file_uploader"))  # type: ignore

    @property
    def expander(self) -> Sequence[Expander]:
        return self.get("expander")  # type: ignore

    @property
    def header(self) -> ElementList[Header]:
        return ElementList(self.get("header"))  # type: ignore

    @property
    def image(self) -> ElementList[Image]:
        return ElementList(self.get("image"))  # type: ignore

    @property
    def info(self) -> ElementList[Info]:
        return ElementList(self.get("info"))  # type: ignore

    @property
    def json(self) -> ElementList[Json]:
        return ElementList(self.get("json"))  # type: ignore

    @property
    def latex(self) -> ElementList[Latex]:
        return ElementList(self.get("latex"))  # type: ignore

    @property
    def markdown(self) -> ElementList[Markdown]:
        return ElementList(self.get("markdown"))  # type: ignore

    @property
    def metric(self) -> ElementList[Metric]:
        return ElementList(self.get("metric"))  # type: ignore

    @property
    def menu_button(self) -> WidgetList[MenuButton[Any]]:
        return WidgetList(self.get("menu_button"))  # type: ignore

    @property
    def multiselect(self) -> WidgetList[Multiselect[Any]]:
        return WidgetList(self.get("multiselect"))  # type: ignore

    @property
    def number_input(self) -> WidgetList[NumberInput]:
        return WidgetList(self.get("number_input"))  # type: ignore

    @property
    def radio(self) -> WidgetList[Radio[Any]]:
        return WidgetList(self.get("radio"))  # type: ignore

    @property
    def select_slider(self) -> WidgetList[SelectSlider[Any]]:
        return WidgetList(self.get("select_slider"))  # type: ignore

    @property
    def selectbox(self) -> WidgetList[Selectbox[Any]]:
        return WidgetList(self.get("selectbox"))  # type: ignore

    @property
    def slider(self) -> WidgetList[Slider[Any]]:
        return WidgetList(self.get("slider"))  # type: ignore

    @property
    def space(self) -> ElementList[Space]:
        return ElementList(self.get("space"))  # type: ignore

    @property
    def status(self) -> Sequence[Status]:
        return self.get("status")  # type: ignore

    @property
    def subheader(self) -> ElementList[Subheader]:
        return ElementList(self.get("subheader"))  # type: ignore

    @property
    def success(self) -> ElementList[Success]:
        return ElementList(self.get("success"))  # type: ignore

    @property
    def table(self) -> ElementList[Table]:
        return ElementList(self.get("table"))  # type: ignore

    @property
    def tabs(self) -> Sequence[Tab]:
        return self.get("tab")  # type: ignore

    @property
    def text(self) -> ElementList[Text]:
        return ElementList(self.get("text"))  # type: ignore

    @property
    def text_area(self) -> WidgetList[TextArea]:
        return WidgetList(self.get("text_area"))  # type: ignore

    @property
    def text_input(self) -> WidgetList[TextInput]:
        return WidgetList(self.get("text_input"))  # type: ignore

    @property
    def time_input(self) -> WidgetList[TimeInput]:
        return WidgetList(self.get("time_input"))  # type: ignore

    @property
    def title(self) -> ElementList[Title]:
        return ElementList(self.get("title"))  # type: ignore

    @property
    def toast(self) -> ElementList[Toast]:
        return ElementList(self.get("toast"))  # type: ignore

    @property
    def toggle(self) -> WidgetList[Toggle]:
        return WidgetList(self.get("toggle"))  # type: ignore

    @property
    def warning(self) -> ElementList[Warning]:
        return ElementList(self.get("warning"))  # type: ignore

    def get(self, element_type: str) -> Sequence[Node]:
        """Return nodes for an AppTest collection name or a node type.

        Public names that differ from ``Node.type`` (for example
        ``datetime_input`` vs ``date_time_input``) are accepted. Node type
        names (usually the proto field name) keep working. ``pills`` /
        ``segmented_control`` / ``container`` use the same filtering as the
        matching attributes.
        """
        if element_type == "pills":
            return list(self.pills)
        if element_type == "segmented_control":
            return list(self.segmented_control)
        if element_type == "container":
            return list(self.container)
        resolved = _GET_TYPE_ALIASES.get(element_type, element_type)
        return [e for e in self if e.type == resolved]

    def run(self, *, timeout: float | None = None) -> AppTest:
        """Run the script with updated widget values.

        Parameters
        ----------
        timeout
            The maximum number of seconds to run the script. None means
            use the AppTest's default.
        """
        return self.root.run(timeout=timeout)

    def __repr__(self) -> str:
        return repr_(self)


def repr_(self: object) -> str:
    """A custom repr similar to `streamlit.util.repr_` but that shows tree
    structure using indentation.
    """
    classname = self.__class__.__name__

    defaults: list[Any] = [None, "", False, [], set(), {}]

    if is_dataclass(self):
        fields_vals = (
            (f.name, getattr(self, f.name))
            for f in fields(self)
            if f.repr
            and getattr(self, f.name) != f.default
            and getattr(self, f.name) not in defaults
        )
    else:
        fields_vals = ((f, v) for (f, v) in self.__dict__.items() if v not in defaults)

    reprs = []
    for field_name, value in fields_vals:
        line = (
            f"{field_name}={format_dict(value)}"
            if isinstance(value, dict)
            else f"{field_name}={value!r}"
        )
        reprs.append(line)

    reprs[0] = "\n" + reprs[0]
    field_reprs = ",\n".join(reprs)

    field_reprs = textwrap.indent(field_reprs, " " * 4)
    return f"{classname}({field_reprs}\n)"


def format_dict(d: dict[Any, Any]) -> str:
    lines = []
    for k, v in d.items():
        line = f"{k}: {v!r}"
        lines.append(line)
    r = ",\n".join(lines)
    r = textwrap.indent(r, " " * 4)
    return f"{{\n{r}\n}}"


@dataclass(repr=False)
class SpecialBlock(Block):
    """Base class for the sidebar and main body containers."""

    def __init__(
        self,
        proto: BlockProto | None,
        root: ElementTree,
        type: str | None = None,
    ) -> None:
        self.children = {}
        self.proto = proto
        if type:
            self.type = type
        elif proto and proto.WhichOneof("type"):
            ty = proto.WhichOneof("type")
            assert ty is not None
            self.type = ty
        else:
            self.type = "unknown"
        self.root = root
        self._block_id = getattr(proto, "id", "") or "" if proto is not None else ""


@dataclass(repr=False)
class ChatMessage(Block):
    """A representation of ``st.chat_message``."""

    type: str = field(repr=False)
    proto: BlockProto.ChatMessage = field(repr=False)
    name: str
    avatar: str

    def __init__(
        self,
        proto: BlockProto.ChatMessage,
        root: ElementTree,
    ) -> None:
        self.children = {}
        self.proto = proto
        self.root = root
        self.type = "chat_message"
        self.name = proto.name
        self.avatar = proto.avatar


@dataclass(repr=False)
class Column(Block):
    """A representation of a column within ``st.columns``."""

    type: str = field(repr=False)
    proto: BlockProto.Column = field(repr=False)
    weight: float
    gap: str | None

    # Mapping from GapSize enum to string
    _GAP_SIZE_TO_STRING: ClassVar[dict[int, str | None]] = {
        GapSize.NONE: None,
        GapSize.XXSMALL: "xxsmall",
        GapSize.XSMALL: "xsmall",
        GapSize.SMALL: "small",
        GapSize.MEDIUM: "medium",
        GapSize.LARGE: "large",
        GapSize.XLARGE: "xlarge",
        GapSize.XXLARGE: "xxlarge",
    }

    def __init__(
        self,
        proto: BlockProto.Column,
        root: ElementTree,
    ) -> None:
        self.children = {}
        self.proto = proto
        self.root = root
        self.type = "column"
        self.weight = proto.weight
        self.gap = self._GAP_SIZE_TO_STRING.get(proto.gap_config.gap_size)


@dataclass(repr=False)
class Expander(Block):
    type: str = field(repr=False)
    proto: BlockProto.Expandable = field(repr=False)
    icon: str
    label: str

    def __init__(self, proto: BlockProto.Expandable, root: ElementTree) -> None:
        self.children = {}
        self.proto = proto
        self.root = root
        # The internal name is "expandable" but the public API uses "expander"
        # so the naming of the class and type follows the public name.
        self.type = "expander"
        self.icon = proto.icon
        self.label = proto.label


@dataclass(repr=False)
class Status(Block):
    type: str = field(repr=False)
    proto: BlockProto.Expandable = field(repr=False)
    icon: str
    label: str

    def __init__(self, proto: BlockProto.Expandable, root: ElementTree) -> None:
        self.children = {}
        self.proto = proto
        self.root = root
        self.type = "status"
        self.icon = proto.icon
        self.label = proto.label

    @property
    def state(self) -> str:
        return self.proto.State.Name(self.proto.state).lower()


@dataclass(repr=False)
class Tab(Block):
    """A representation of tab within ``st.tabs``."""

    type: str = field(repr=False)
    proto: BlockProto.Tab = field(repr=False)
    label: str

    def __init__(
        self,
        proto: BlockProto.Tab,
        root: ElementTree,
    ) -> None:
        self.children = {}
        self.proto = proto
        self.root = root
        self.type = "tab"
        self.label = proto.label


Node: TypeAlias = Element | Block


def _widget_form_id(node: Widget) -> str:
    """Return the widget's form id, or ``""`` if it is not in a form."""
    return getattr(node.proto, "form_id", "") or ""


def _submitted_form_ids(tree: ElementTree) -> set[str]:
    """Form ids whose submit button is triggered for this ``.run()``."""
    submitted: set[str] = set()
    for node in tree:
        if isinstance(node, Button) and node._value:
            form_id = _widget_form_id(node)
            if form_id:
                submitted.add(form_id)
    return submitted


def _form_clear_flags(tree: ElementTree) -> dict[str, bool]:
    """Map form id → ``clear_on_submit`` for every form in ``tree``."""
    flags: dict[str, bool] = {}
    for node in tree:
        if getattr(node, "type", None) != "form":
            continue
        proto = getattr(node, "proto", None)
        form = getattr(proto, "form", None)
        if form is not None:
            flags[form.form_id] = bool(form.clear_on_submit)
    return flags


def _unset_value_marker(node: Widget) -> tuple[str, Any]:
    """Attribute name and unset marker so serialization uses the committed value.

    Each widget class uses a different "not staged" marker: ``InitialValue``,
    ``None``, or ``False`` for buttons.
    """
    if isinstance(node, FileUploader):
        return ("_files", InitialValue())
    if isinstance(node, (Button, DownloadButton)):
        return ("_value", False)
    if isinstance(
        node,
        (
            DateInput,
            DateTimeInput,
            Feedback,
            NumberInput,
            Radio,
            Selectbox,
            TextArea,
            TextInput,
            TimeInput,
            ButtonGroup,
        ),
    ):
        return ("_value", InitialValue())
    return ("_value", None)


def _has_pending_value(node: Widget) -> bool:
    """Return True if the test staged a value on this widget since the last run.

    ``None`` is a real staged value for widgets whose unset marker is
    ``InitialValue`` (for example ``selectbox.select_index(None)``).
    """
    attr, sentinel = _unset_value_marker(node)
    current = getattr(node, attr)
    if isinstance(sentinel, InitialValue):
        return not isinstance(current, InitialValue)
    return current is not sentinel


def _use_form_clear_defaults(
    node: Widget,
    *,
    submitted: set[str],
    cleared: set[str],
    form_clears: dict[str, bool],
) -> bool:
    """Return True if this widget should serialize its form's proto default.

    Gating on the form's *current* ``clear_on_submit`` is an AppTest
    approximation: the frontend stages cleared defaults into the form's pending
    widget states even if the next render sets ``clear_on_submit=False``.
    """
    form_id = _widget_form_id(node)
    return bool(
        form_id
        and form_id in submitted
        and form_id in cleared
        and form_clears.get(form_id)
        and not isinstance(node, (Button, DownloadButton))
        and not _has_pending_value(node)
    )


def _record_submitted_form_clears(
    runner: AppTest, submitted: set[str], form_clears: dict[str, bool]
) -> None:
    """Remember which submitted forms should send defaults on the next submit."""
    for form_id in submitted:
        if form_clears.get(form_id):
            runner._cleared_form_ids.add(form_id)
        else:
            runner._cleared_form_ids.discard(form_id)


def _cleared_widget_state(node: Widget) -> WidgetState:
    """``WidgetState`` that deserializes to the widget's declared default.

    An unset value oneof makes ``session_state`` call ``deserializer(None)``,
    which is the canonical default path and does not re-run ``format_func``.
    """
    ws = WidgetState()
    ws.id = node.id
    return ws


def get_widget_state(node: Node) -> WidgetState | None:
    if isinstance(node, Widget):
        return node._widget_state
    return None


@dataclass(repr=False)
class ElementTree(Block):
    """A tree of the elements produced by running a streamlit script.

    Elements can be queried in three ways:
    - By element type, using `.foo` properties to get a list of all of that element,
    in the order they appear in the app
    - By user key, for widgets, by calling the above list with a key: `.foo(key='bar')`
    - Positionally, using list indexing syntax (`[...]`) to access a child of a
    block element. Not recommended because the exact tree structure can be surprising.

    Element queries made on a block container will return only the elements
    descending from that block.

    Returned elements have methods for accessing whatever attributes are relevant.
    For very simple elements this may be only its value, while complex elements
    like widgets have many.

    Widgets provide a fluent API for faking frontend interaction and rerunning
    the script with the new widget values. All widgets provide a low level `set_value`
    method, along with higher level methods specific to that type of widget.
    After an interaction, calling `.run()` will update the AppTest with the
    results of that script run.
    """

    _runner: AppTest | None = field(repr=False, default=None)

    def __init__(self) -> None:
        self.children = {}
        self.root = self
        self.type = "root"
        self._block_id = ""

    @property
    def main(self) -> Block:
        m = self[0]
        assert isinstance(m, Block)
        return m

    @property
    def sidebar(self) -> Block:
        s = self[1]
        assert isinstance(s, Block)
        return s

    @property
    def session_state(self) -> SafeSessionState:
        assert self._runner is not None
        # Widget internals need SafeSessionState; AppTest.session_state is the
        # dict-like wrapper testers use.
        return self._runner._session_state

    def get_widget_states(self) -> WidgetStates:
        """Serialize widget values for the next script run.

        Form widgets are included so a new ScriptRunner does not cull them, but
        uncommitted ``set_value`` / ``click`` is ignored until that form's
        submit button is triggered. After ``clear_on_submit``, the next submit
        serializes proto defaults for widgets the test has not set again.
        """
        submitted = _submitted_form_ids(self)
        form_clears = _form_clear_flags(self)
        runner = self._runner
        cleared: set[str] = runner._cleared_form_ids if runner is not None else set()

        ws = WidgetStates()
        for node in self:
            if not isinstance(node, Widget):
                continue
            form_id = _widget_form_id(node)

            if _use_form_clear_defaults(
                node,
                submitted=submitted,
                cleared=cleared,
                form_clears=form_clears,
            ):
                ws.widgets.append(_cleared_widget_state(node))
                continue

            restore: tuple[str, Any] | None = None
            if form_id and form_id not in submitted:
                attr, sentinel = _unset_value_marker(node)
                restore = (attr, getattr(node, attr))
                setattr(node, attr, sentinel)
            try:
                w = get_widget_state(node)
            finally:
                if restore is not None:
                    setattr(node, restore[0], restore[1])
            if w is not None:
                ws.widgets.append(w)

        return ws

    def run(self, *, timeout: float | None = None) -> AppTest:
        """Run the script with updated widget values.

        Parameters
        ----------
        timeout
            The maximum number of seconds to run the script. None means
            use the AppTest's default.
        """
        assert self._runner is not None

        submitted = _submitted_form_ids(self)
        form_clears = _form_clear_flags(self)
        widget_states = self.get_widget_states()
        result = self._runner._run(widget_states, timeout=timeout)
        _record_submitted_form_clears(self._runner, submitted, form_clears)
        return result

    def __repr__(self) -> str:
        return format_dict(self.children)


def parse_tree_from_messages(messages: list[ForwardMsg]) -> ElementTree:
    """Transform a list of `ForwardMsg` into a tree matching the implicit
    tree structure of blocks and elements in a streamlit app.

    Returns the root of the tree, which acts as the entrypoint for the query
    and interaction API.
    """
    root = ElementTree()
    root.children = {
        0: SpecialBlock(type="main", root=root, proto=None),
        1: SpecialBlock(type="sidebar", root=root, proto=None),
        2: SpecialBlock(type="event", root=root, proto=None),
    }

    for msg in messages:
        if not msg.HasField("delta"):
            continue
        delta_path = msg.metadata.delta_path
        delta = msg.delta
        if delta.WhichOneof("type") == "new_element":
            elt = delta.new_element
            ty = elt.WhichOneof("type")
            new_node: Node
            if ty == "alert":
                alert_format = elt.alert.format
                if alert_format == AlertProto.Format.ERROR:
                    new_node = Error(elt.alert, root=root)
                elif alert_format == AlertProto.Format.INFO:
                    new_node = Info(elt.alert, root=root)
                elif alert_format == AlertProto.Format.SUCCESS:
                    new_node = Success(elt.alert, root=root)
                elif alert_format == AlertProto.Format.WARNING:
                    new_node = Warning(elt.alert, root=root)
                else:
                    new_node = UnknownElement(elt, root=root)
            elif ty == "dataframe":
                new_node = Dataframe(elt.dataframe, root=root)
            elif ty == "table":
                new_node = Table(elt.table, root=root)
            elif ty == "button":
                new_node = Button(elt.button, root=root)
            elif ty == "button_group":
                new_node = ButtonGroup(elt.button_group, root=root)
            elif ty == "chat_input":
                new_node = ChatInput(elt.chat_input, root=root)
            elif ty == "checkbox":
                style = elt.checkbox.type
                if style == CheckboxProto.StyleType.TOGGLE:
                    new_node = Toggle(elt.checkbox, root=root)
                else:
                    new_node = Checkbox(elt.checkbox, root=root)
            elif ty == "code":
                new_node = Code(elt.code, root=root)
            elif ty == "color_picker":
                new_node = ColorPicker(elt.color_picker, root=root)
            elif ty == "date_input":
                new_node = DateInput(elt.date_input, root=root)
            elif ty == "date_time_input":
                new_node = DateTimeInput(elt.date_time_input, root=root)
            elif ty == "download_button":
                new_node = DownloadButton(elt.download_button, root=root)
            elif ty == "exception":
                new_node = Exception(elt.exception, root=root)
            elif ty == "feedback":
                new_node = Feedback(elt.feedback, root=root)
            elif ty == "file_uploader":
                new_node = FileUploader(elt.file_uploader, root=root)
            elif ty == "heading":
                if elt.heading.tag == HeadingProtoTag.TITLE_TAG.value:
                    new_node = Title(elt.heading, root=root)
                elif elt.heading.tag == HeadingProtoTag.HEADER_TAG.value:
                    new_node = Header(elt.heading, root=root)
                elif elt.heading.tag == HeadingProtoTag.SUBHEADER_TAG.value:
                    new_node = Subheader(elt.heading, root=root)
                else:
                    new_node = UnknownElement(elt, root=root)
            elif ty == "imgs":
                new_node = Image(elt.imgs, root=root)
            elif ty == "json":
                new_node = Json(elt.json, root=root)
            elif ty == "markdown":
                if elt.markdown.element_type == MarkdownProto.Type.NATIVE:
                    new_node = Markdown(elt.markdown, root=root)
                elif elt.markdown.element_type == MarkdownProto.Type.CAPTION:
                    new_node = Caption(elt.markdown, root=root)
                elif elt.markdown.element_type == MarkdownProto.Type.LATEX:
                    new_node = Latex(elt.markdown, root=root)
                elif elt.markdown.element_type == MarkdownProto.Type.DIVIDER:
                    new_node = Divider(elt.markdown, root=root)
                else:
                    new_node = UnknownElement(elt, root=root)
            elif ty == "menu_button":
                new_node = MenuButton(elt.menu_button, root=root)
            elif ty == "metric":
                new_node = Metric(elt.metric, root=root)
            elif ty == "multiselect":
                new_node = Multiselect(elt.multiselect, root=root)
            elif ty == "number_input":
                new_node = NumberInput(elt.number_input, root=root)
            elif ty == "radio":
                new_node = Radio(elt.radio, root=root)
            elif ty == "selectbox":
                new_node = Selectbox(elt.selectbox, root=root)
            elif ty == "slider":
                if elt.slider.type == SliderProto.Type.SLIDER:
                    new_node = Slider(elt.slider, root=root)
                elif elt.slider.type == SliderProto.Type.SELECT_SLIDER:
                    new_node = SelectSlider(elt.slider, root=root)
                else:
                    new_node = UnknownElement(elt, root=root)
            elif ty == "space":
                new_node = Space(
                    elt.space,
                    root=root,
                    size=_space_size_from_width_config(elt.width_config),
                )
            elif ty == "text":
                new_node = Text(elt.text, root=root)
            elif ty == "text_area":
                new_node = TextArea(elt.text_area, root=root)
            elif ty == "text_input":
                new_node = TextInput(elt.text_input, root=root)
            elif ty == "time_input":
                new_node = TimeInput(elt.time_input, root=root)
            elif ty == "toast":
                new_node = Toast(elt.toast, root=root)
            else:
                new_node = UnknownElement(elt, root=root)
        elif delta.WhichOneof("type") == "add_block":
            block = delta.add_block
            bty = block.WhichOneof("type")
            if bty == "chat_message":
                new_node = ChatMessage(block.chat_message, root=root)
            elif bty == "column":
                new_node = Column(block.column, root=root)
            elif bty == "expandable":
                # st.status always sets state; st.expander leaves it undefined
                # even when an icon is present.
                if block.expandable.state != block.expandable.State.STATE_UNDEFINED:
                    new_node = Status(block.expandable, root=root)
                else:
                    new_node = Expander(block.expandable, root=root)
            elif bty == "tab":
                new_node = Tab(block.tab, root=root)
            else:
                new_node = Block(proto=block, root=root)
            new_node._block_id = block.id or ""
        elif delta.WhichOneof("type") == "new_transient":
            # new_transient (e.g. spinner) - skip these in the element tree
            continue
        else:
            # Unknown delta type - skip to avoid silently swallowing future types
            continue

        current_node: Block = root
        # Every node up to the end is a Block
        for idx in delta_path[:-1]:
            children = current_node.children
            child = children.get(idx)
            if child is None:
                child = Block(proto=None, root=root)
                children[idx] = child
            assert isinstance(child, Block)
            current_node = child

        # Handle a block when we already have a placeholder for that location
        if isinstance(new_node, Block):
            placeholder_block = current_node.children.get(delta_path[-1])
            if placeholder_block is not None:
                new_node.children = placeholder_block.children

        current_node.children[delta_path[-1]] = new_node

    return root