# 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,
    Generic,
    TypeAlias,
    TypeVar,
    cast,
    overload,
)

from typing_extensions import Self

from streamlit import dataframe_util, util
from streamlit.elements.heading import HeadingProtoTag
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

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.runtime.state.safe_session_state import SafeSessionState
    from streamlit.testing.v1.app_test import AppTest

T = TypeVar("T")


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: ElementProto, root: ElementTree) -> None: ...

    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:
        """Fallback attempt to get an attribute from the proto."""
        return getattr(self.proto, name)

    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
        self.key = None

    @property
    def value(self) -> Any:
        try:
            state = self.root.session_state
            assert state is not None
            return state[self.proto.id]
        except ValueError:
            # No id field, not a widget
            return self.proto.value


@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 set_value(self, v: Any) -> Self:
        """Set the value of the widget."""
        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))

    @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]):
    def __call__(self, key: str) -> W_co:
        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."""
        self._value = v
        return self

    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."""
        self._value = v
        return self

    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."""
        self._value = v
        return self

    @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 | None:
        """The value of the widget. (str)"""  # noqa: D400
        if self._value:
            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."""
        self._value = v
        return self

    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."""
        self._value = v
        return self

    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."""
        self._value = v
        return self

    @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


@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 for testing."""

    proto: SpaceProto = field(repr=False)

    key: None = None

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


@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] | 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.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 self._value is not None:
            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 self._is_single_select:
            # Single-select: value is a single item or None
            if value is None:
                return []
            return [_format_value_for_widget(format_func, value)]
        # Multi-select: value is a list
        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).
        """
        self._value = v
        return self

    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)
        # Multi-select: add to list
        current = cast("list[T]", self.value)
        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
        # Multi-select: remove from list
        current = cast("list[T]", self.value)
        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
        self._value = v
        return self


@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

        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

        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._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."""
        self._value = v
        return self

    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

        self._value = v
        return self

    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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self

    @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."""
        self._value = v
        return self


@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)

    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

    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:
        return None

    # 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 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 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 [e for e in self if e.type == element_type]

    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


@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:
        if self.icon == "spinner":
            return "running"
        if self.icon == ":material/check:":
            return "complete"
        if self.icon == ":material/error:":
            return "error"
        raise ValueError("Unknown Status state")


@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 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"

    @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
        return self._runner.session_state

    def get_widget_states(self) -> WidgetStates:
        ws = WidgetStates()
        for node in self:
            w = get_widget_state(node)
            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

        widget_states = self.get_widget_states()
        return self._runner._run(widget_states, timeout=timeout)

    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:
                    raise ValueError(
                        f"Unknown alert type with format {elt.alert.format}"
                    )
            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:
                    raise ValueError(f"Unknown heading type with tag {elt.heading.tag}")
            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:
                    raise ValueError(
                        f"Unknown markdown type {elt.markdown.element_type}"
                    )
            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:
                    raise ValueError(f"Slider with unknown type {elt.slider}")
            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":
                if block.expandable.icon:
                    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)
        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
