# 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.

# Keep Attributes before Examples in API docstrings.

from __future__ import annotations

import json
import weakref
from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    Literal,
    TypeAlias,
    cast,
    overload,
)

from streamlit.deprecation_util import (
    make_deprecated_name_warning,
    show_deprecation_warning,
)
from streamlit.elements.lib.form_utils import current_form_id
from streamlit.elements.lib.layout_utils import (
    HeightWithoutContent,
    WidthWithoutContent,
    create_layout_config,
)
from streamlit.elements.lib.policies import check_widget_policies
from streamlit.elements.lib.utils import (
    Key,
    compute_and_register_element_id,
    normalize_alt,
    to_key,
)
from streamlit.errors import (
    StreamlitIncompatibleParametersError,
    StreamlitInvalidParameterTypeError,
    StreamlitValueError,
)
from streamlit.proto.DeckGlJsonChart_pb2 import DeckGlJsonChart as PydeckProto
from streamlit.runtime.metrics_util import gather_metrics
from streamlit.runtime.scriptrunner_utils.script_run_context import get_script_run_ctx
from streamlit.runtime.state import (
    WidgetCallback,
    register_widget,
    validate_on_change_mode,
)
from streamlit.util import ReadOnlyAttributeDictionary

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping

    from pydeck import Deck

    from streamlit.delta_generator import DeltaGenerator


# Mapping used when no data is passed.
EMPTY_MAP: Final[Mapping[str, Any]] = {
    "initialViewState": {"latitude": 0, "longitude": 0, "pitch": 0, "zoom": 1},
}

SelectionMode: TypeAlias = Literal["single-object", "multi-object"]
_SELECTION_MODES: Final[set[SelectionMode]] = {
    "single-object",
    "multi-object",
}


def parse_selection_mode(
    selection_mode: SelectionMode | Iterable[SelectionMode],
) -> set[PydeckProto.SelectionMode.ValueType]:
    """Parse and check the user provided selection modes."""
    if isinstance(selection_mode, str):
        # Only a single selection mode was passed
        selection_mode_set = {selection_mode}
    else:
        # Only a single string selection mode is supported. Lists and sets are
        # rejected until multi-mode selection lands (see DeckGlJsonChart.tsx).
        raise StreamlitInvalidParameterTypeError(
            "selection_mode",
            type(selection_mode).__name__,
            ["str"],
            detail="Selection mode must be a single value.",
        )

    if not selection_mode_set.issubset(_SELECTION_MODES):
        raise StreamlitValueError(
            "selection_mode",
            [f"'{mode}'" for mode in sorted(_SELECTION_MODES)],
        )

    if selection_mode_set.issuperset(  # pragma: no cover - defensive, only string inputs reach here
        {"single-object", "multi-object"}
    ):
        raise StreamlitIncompatibleParametersError(
            "selection_mode='single-object'",
            "selection_mode='multi-object'",
        )

    parsed_selection_modes = []
    for mode in selection_mode_set:
        if mode == "single-object":
            parsed_selection_modes.append(PydeckProto.SelectionMode.SINGLE_OBJECT)
        elif mode == "multi-object":
            parsed_selection_modes.append(PydeckProto.SelectionMode.MULTI_OBJECT)
    return set(parsed_selection_modes)


class PydeckSelectionState(ReadOnlyAttributeDictionary):
    r"""
    The schema for the PyDeck chart selection state.

    The selection state is stored in a read-only dictionary-like object that
    supports both key and attribute notation. Selection states cannot be
    programmatically changed or set through Session State.

    You must define ``id`` in ``pydeck.Layer`` to ensure statefulness when
    using selections with ``st.pydeck_chart``.

    Attributes
    ----------
    indices : dict[str, list[int]]
        A dictionary of selected objects by layer. Each key in the dictionary
        is a layer id, and each value is a list of object indices within that
        layer.
    objects : dict[str, list[dict[str, Any]]]
        A dictionary of object attributes by layer. Each key in the dictionary
        is a layer id, and each value is a list of metadata dictionaries for
        the selected objects in that layer.

    Examples
    --------
    The following example has multi-object selection enabled. The chart
    displays US state capitals by population (2023 US Census estimate). You
    can access this `data
    <https://github.com/streamlit/docs/blob/main/python/api-examples-source/data/capitals.csv>`_
    from GitHub.

    >>> import streamlit as st
    >>> import pydeck
    >>> import pandas as pd
    >>>
    >>> capitals = pd.read_csv(
    ...     "capitals.csv",
    ...     header=0,
    ...     names=[
    ...         "Capital",
    ...         "State",
    ...         "Abbreviation",
    ...         "Latitude",
    ...         "Longitude",
    ...         "Population",
    ...     ],
    ... )
    >>> capitals["size"] = capitals.Population / 10
    >>>
    >>> point_layer = pydeck.Layer(
    ...     "ScatterplotLayer",
    ...     data=capitals,
    ...     id="capital-cities",
    ...     get_position=["Longitude", "Latitude"],
    ...     get_color="[255, 75, 75]",
    ...     pickable=True,
    ...     auto_highlight=True,
    ...     get_radius="size",
    ... )
    >>>
    >>> view_state = pydeck.ViewState(
    ...     latitude=40, longitude=-117, controller=True, zoom=2.4, pitch=30
    ... )
    >>>
    >>> chart = pydeck.Deck(
    ...     point_layer,
    ...     initial_view_state=view_state,
    ...     tooltip={"text": "{Capital}, {Abbreviation}\nPopulation: {Population}"},
    ... )
    >>>
    >>> event = st.pydeck_chart(chart, on_select="rerun", selection_mode="multi-object")
    >>>
    >>> event.selection

    .. output::
        https://doc-pydeck-event-state-selections.streamlit.app/
        height: 700px

    This is an example of the selection state when selecting a single object
    from a layer with id, ``"captial-cities"``:

    >>> {
    >>>   "indices":{
    >>>     "capital-cities":[
    >>>       2
    >>>     ]
    >>>   },
    >>>   "objects":{
    >>>     "capital-cities":[
    >>>       {
    >>>         "Abbreviation":" AZ"
    >>>         "Capital":"Phoenix"
    >>>         "Latitude":33.448457
    >>>         "Longitude":-112.073844
    >>>         "Population":1650070
    >>>         "State":" Arizona"
    >>>         "size":165007.0
    >>>       }
    >>>     ]
    >>>   }
    >>> }

    """

    indices: dict[str, list[int]]
    objects: dict[str, list[dict[str, Any]]]

    @overload
    def __getitem__(self, key: Literal["indices"]) -> dict[str, list[int]]: ...

    @overload
    def __getitem__(
        self, key: Literal["objects"]
    ) -> dict[str, list[dict[str, Any]]]: ...

    @overload
    def __getitem__(self, key: Any) -> Any: ...

    def __getitem__(self, key: Any) -> Any:
        return super().__getitem__(key)


class PydeckState(ReadOnlyAttributeDictionary):
    """
    The schema for the PyDeck event state.

    To use this type in an annotation, import it from ``streamlit.typing``.

    The event state is stored in a read-only dictionary-like object that
    supports both key and attribute notation. Event states cannot be
    programmatically changed or set through Session State.

    Only selection events are supported at this time.

    Attributes
    ----------
    selection : dict
        The state of the ``on_select`` event. This attribute returns a
        dictionary-like object that supports both key and attribute notation.
        The attributes are described by ``PydeckSelectionState``.

    """

    selection: PydeckSelectionState

    # ReadOnlyAttributeDictionary routes attribute access through __getitem__,
    # so the override below is enough to keep `selection` typed as
    # PydeckSelectionState. Use dict.__getitem__ for the selection key so the
    # read-only base class does not re-wrap the already-typed nested instance.
    @overload
    def __getitem__(self, key: Literal["selection"]) -> PydeckSelectionState: ...

    @overload
    def __getitem__(self, key: Any) -> Any: ...

    def __getitem__(self, key: Any) -> Any:
        if key == "selection":
            item = dict.__getitem__(self, key)
            if not isinstance(item, PydeckSelectionState):
                item = PydeckSelectionState(item)
                # Cache so repeated bracket/attribute access stays identity-stable.
                dict.__setitem__(self, key, item)
            return item
        return super().__getitem__(key)


@dataclass
class PydeckSelectionSerde:
    """PydeckSelectionSerde is used to serialize and deserialize the Pydeck selection state."""

    def deserialize(self, ui_value: str | None) -> PydeckState:
        empty_selection_state = PydeckState(
            {
                "selection": PydeckSelectionState(
                    {
                        "indices": {},
                        "objects": {},
                    }
                )
            }
        )

        if ui_value is None:
            return empty_selection_state

        selection_state = json.loads(ui_value)
        # We have seen some situations where the ui_value was just an empty
        # dict, so we want to ensure that it always returns the empty state in
        # case this happens.
        if "selection" not in selection_state:
            return empty_selection_state

        # Eagerly wrap selection so bracket access returns a stable typed
        # instance instead of creating a shallow copy on every access.
        selection_state["selection"] = PydeckSelectionState(
            selection_state["selection"]
        )
        return PydeckState(selection_state)

    def serialize(self, selection_state: PydeckState) -> str:
        return json.dumps(selection_state, default=str)


class PydeckMixin:
    @overload
    def pydeck_chart(
        self,
        pydeck_obj: Deck | None = None,
        *,
        width: WidthWithoutContent = "stretch",
        use_container_width: bool | None = None,
        height: HeightWithoutContent = 500,
        selection_mode: SelectionMode = "single-object",
        on_select: Literal["ignore"] = "ignore",
        key: Key | None = None,
        alt: str | None = None,
    ) -> DeltaGenerator: ...

    @overload
    def pydeck_chart(
        self,
        pydeck_obj: Deck | None = None,
        *,
        width: WidthWithoutContent = "stretch",
        use_container_width: bool | None = None,
        height: HeightWithoutContent = 500,
        selection_mode: SelectionMode = "single-object",
        # No default: omitted on_select must match the "ignore" overload.
        on_select: Literal["rerun"] | WidgetCallback,
        key: Key | None = None,
        alt: str | None = None,
    ) -> PydeckState: ...

    @gather_metrics("pydeck_chart")
    def pydeck_chart(
        self,
        pydeck_obj: Deck | None = None,
        *,
        width: WidthWithoutContent = "stretch",
        use_container_width: bool | None = None,
        height: HeightWithoutContent = 500,
        selection_mode: SelectionMode = "single-object",
        on_select: Literal["rerun", "ignore"] | WidgetCallback = "ignore",
        key: Key | None = None,
        alt: str | None = None,
    ) -> DeltaGenerator | PydeckState:
        """Draw a chart using the PyDeck library.

        This supports 3D maps, point clouds, and more! More info about PyDeck
        at https://deckgl.readthedocs.io/en/latest/.

        These docs are also quite useful:

        - DeckGL docs: https://github.com/uber/deck.gl/tree/master/docs
        - DeckGL JSON docs: https://github.com/uber/deck.gl/tree/master/modules/json

        When using this command, a service called Carto_ provides the map tiles to render
        map content. If you're using advanced PyDeck features you may need to obtain
        an API key from Carto first. You can do that as
        ``pydeck.Deck(api_keys={"carto": YOUR_KEY})`` or by setting the CARTO_API_KEY
        environment variable. See `PyDeck's documentation`_ for more information.

        Another common provider for map tiles is Mapbox_. If you prefer to use that,
        you'll need to create an account at https://mapbox.com and specify your Mapbox
        key when creating the ``pydeck.Deck`` object. You can do that as
        ``pydeck.Deck(api_keys={"mapbox": YOUR_KEY})`` or by setting the MAPBOX_API_KEY
        environment variable.

        .. _Carto: https://carto.com
        .. _Mapbox: https://mapbox.com
        .. _PyDeck's documentation: https://deckgl.readthedocs.io/en/latest/deck.html

        Carto and Mapbox are third-party products and Streamlit accepts no responsibility
        or liability of any kind for Carto or Mapbox, or for any content or information
        made available by Carto or Mapbox. The use of Carto or Mapbox is governed by
        their respective Terms of Use.

        .. note::
            Pydeck uses two WebGL contexts per chart, and different browsers
            have different limits on the number of WebGL contexts per page.
            If you exceed this limit, the oldest contexts will be dropped to
            make room for the new ones. To avoid this limitation in most
            browsers, don't display more than eight Pydeck charts on a single
            page.

        Parameters
        ----------
        pydeck_obj : pydeck.Deck or None
            Object specifying the PyDeck chart to draw. Built-in deck.gl
            layers, views (``MapView``, ``OrbitView``, ``OrthographicView``,
            ``FirstPersonView``, and ``GlobeView``), JSON ``parameters``, and
            layer extensions are supported. Use ``map_provider=None`` to omit
            the basemap (pydeck 0.9+). Pass extensions as ``@@type`` dicts,
            for example
            ``extensions=[{"@@type": "DataFilterExtension", "filterSize": 1}]``.
            Custom JS libraries, widgets, and multi-view layouts are not
            supported.
        width : "stretch" or int
            The width of the chart element. This can be one of the following:

            - ``"stretch"`` (default): The width of the element matches the
              width of the parent container.
            - An integer specifying the width in pixels: The element has a
              fixed width. If the specified width is greater than the width of
              the parent container, the width of the element matches the width
              of the parent container.

        use_container_width : bool or None
            Whether to override the chart's native width with the width of
            the parent container. This can be one of the following:

            - ``None`` (default): Streamlit will use the chart's default behavior.
            - ``True``: Streamlit sets the width of the chart to match the
              width of the parent container.
            - ``False``: Streamlit sets the width of the chart to fit its
              contents according to the plotting library, up to the width of
              the parent container.

            .. deprecated::
               ``use_container_width`` is deprecated and will be removed in a
                future release. For ``use_container_width=True``, use
                ``width="stretch"``.

        height : "stretch" or int
            The height of the chart element. This can be one of the following:

            - An integer specifying the height in pixels: The element has a
              fixed height. If the content is larger than the specified
              height, scrolling is enabled. This is ``500`` by default.
            - ``"stretch"``: The height of the element matches the height of
              its content or the height of the parent container, whichever is
              larger. If the element is not in a parent container, the height
              of the element matches the height of its content.

        on_select : "ignore" or "rerun" or callable
            How the figure should respond to user selection events. This controls
            whether or not the chart behaves like an input widget.
            ``on_select`` can be one of the following:

            - ``"ignore"`` (default): Streamlit will not react to any selection
              events in the chart. The figure will not behave like an
              input widget.
            - ``"rerun"``: Streamlit will rerun the app when the user selects
              data in the chart. In this case, ``st.pydeck_chart`` will return
              the selection data as a dictionary.
            - A ``callable``: Streamlit will rerun the app and execute the callable
              as a callback function before the rest of the app. In this case,
              ``st.pydeck_chart`` will return the selection data as a
              dictionary.

            If ``on_select`` is not ``"ignore"``, all layers must have a
            declared ``id`` to keep the chart stateful across reruns.
        selection_mode : "single-object" or "multi-object"
            The selection mode of the chart. This can be one of the following:

            - ``"single-object"`` (default): Only one object can be selected at
              a time.
            - ``"multi-object"``: Multiple objects can be selected at a time.

        key : str, int, or None
            An optional string to use for giving this element a stable
            identity. If this is ``None`` (default), the element's identity
            will be determined based on the values of the other parameters.

            If selections are activated and ``key`` is provided,
            Streamlit will register the key in Session State to store the
            selection state. The selection state is read-only. For more
            details, see `Widget behavior
            <https://docs.streamlit.io/develop/concepts/architecture/widget-behavior>`_.

            Additionally, if ``key`` is provided, it will be used as a
            CSS class name prefixed with ``st-key-``.

        alt : str or None
            A description of the chart for screen readers and other assistive
            technologies. If this is ``None`` (default), Streamlit does not
            provide an accessible name for the chart.

            An empty or whitespace-only string is treated the same as ``None``
            and is logged so authors notice the dual meaning of ``alt=""``
            across commands (decorative only on ``st.image`` / ``st.pyplot``).

            Keep this to a short description of the visual; it is not a full
            text alternative for dense graphics.

        Returns
        -------
        element or PydeckState
            If ``on_select`` is ``"ignore"`` (default), this command returns an
            internal placeholder for the chart element. Otherwise, this method
            returns a ``PydeckState`` object. This object is dictionary-like
            and supports both key and attribute notation. To use this type in
            an annotation, import it from ``streamlit.typing``.

        Examples
        --------
        Here's a chart using a HexagonLayer and a ScatterplotLayer. It uses either the
        light or dark map style, based on which Streamlit theme is currently active:

        >>> import pandas as pd
        >>> import pydeck as pdk
        >>> import streamlit as st
        >>> from numpy.random import default_rng as rng
        >>>
        >>> df = pd.DataFrame(
        ...     rng(0).standard_normal((1000, 2)) / [50, 50] + [37.76, -122.4],
        ...     columns=["lat", "lon"],
        ... )
        >>>
        >>> st.pydeck_chart(
        ...     pdk.Deck(
        ...         map_style=None,  # Use Streamlit theme to pick map style
        ...         initial_view_state=pdk.ViewState(
        ...             latitude=37.76,
        ...             longitude=-122.4,
        ...             zoom=11,
        ...             pitch=50,
        ...         ),
        ...         layers=[
        ...             pdk.Layer(
        ...                 "HexagonLayer",
        ...                 data=df,
        ...                 get_position="[lon, lat]",
        ...                 radius=200,
        ...                 elevation_scale=4,
        ...                 elevation_range=[0, 1000],
        ...                 pickable=True,
        ...                 extruded=True,
        ...             ),
        ...             pdk.Layer(
        ...                 "ScatterplotLayer",
        ...                 data=df,
        ...                 get_position="[lon, lat]",
        ...                 get_color="[200, 30, 0, 160]",
        ...                 get_radius=200,
        ...             ),
        ...         ],
        ...     ),
        ...     alt="Hexagon and scatter map of sample points near San Francisco",
        ... )

        .. output::
           https://doc-pydeck-chart.streamlit.app/
           height: 530px

        .. note::
           To make the PyDeck chart's style consistent with Streamlit's theme,
           you can set ``map_style=None`` in the ``pydeck.Deck`` object.

        """
        if use_container_width is not None:
            show_deprecation_warning(
                make_deprecated_name_warning(
                    "use_container_width",
                    "width",
                    "2025-12-31",
                    "For `use_container_width=True`, use `width='stretch'`. "
                    "For `use_container_width=False`, specify an integer width.",
                    include_st_prefix=False,
                ),
                show_in_browser=False,
            )
            if use_container_width:
                width = "stretch"
            # Otherwise keep the provided width.

        layout_config = create_layout_config(width=width, height=height)

        pydeck_proto = PydeckProto()

        ctx = get_script_run_ctx()

        # Workaround for pandas 3.x compatibility issue in pydeck's serialization.
        # See: https://github.com/visgl/deck.gl/issues/9986
        from streamlit.dataframe_util import is_pandas_version_less_than

        if not is_pandas_version_less_than("3.0.0"):
            _prepare_pydeck_for_json(pydeck_obj)

        spec = json.dumps(EMPTY_MAP) if pydeck_obj is None else pydeck_obj.to_json()

        pydeck_proto.json = spec

        tooltip = _get_pydeck_tooltip(pydeck_obj)
        if tooltip:
            pydeck_proto.tooltip = json.dumps(tooltip)

        mapbox_token = getattr(pydeck_obj, "mapbox_key", None)
        if mapbox_token:
            pydeck_proto.mapbox_token = mapbox_token

        normalized_alt = normalize_alt(alt)
        if normalized_alt is not None:
            # Carry alt on its own proto field. The pydeck JSON is hashed into the
            # element ID, and alt must never be written into a hashed spec.
            pydeck_proto.alt = normalized_alt

        key = to_key(key)
        is_selection_activated = on_select != "ignore"

        on_select_callback = validate_on_change_mode(
            on_select,
            supported_modes=("rerun", "ignore"),
            none_supported=False,
            param_name="on_select",
        )

        if is_selection_activated:
            # Selections are activated, treat Pydeck as a widget:
            pydeck_proto.selection_mode.extend(parse_selection_mode(selection_mode))

            # Run some checks that are only relevant when selections are activated
            is_callback = on_select_callback is not None
            check_widget_policies(
                self.dg,
                key,
                on_change=on_select_callback,
                default_value=None,
                writes_allowed=False,
                enable_check_callback_rules=is_callback,
            )
            pydeck_proto.form_id = current_form_id(self.dg)

            pydeck_proto.id = compute_and_register_element_id(
                "deck_gl_json_chart",
                user_key=key,
                # Only selection_mode is hashed when a key is given, so alt changes never
                # reset selection state; unkeyed charts hash alt like any other stable kwarg.
                # Spec/data are also omitted from that keyed hash, which can leave orphaned
                # selections if data length shrinks; the frontend sanitizes invalid indices.
                key_as_main_identity={"selection_mode"},
                dg=self.dg,
                is_selection_activated=is_selection_activated,
                selection_mode=selection_mode,
                use_container_width=use_container_width,
                spec=spec,
                alt=normalized_alt,
            )

            serde = PydeckSelectionSerde()

            widget_state = register_widget(
                pydeck_proto.id,
                ctx=ctx,
                deserializer=serde.deserialize,
                on_change_handler=on_select_callback,
                serializer=serde.serialize,
                value_type="string_value",
            )

            self.dg._enqueue(
                "deck_gl_json_chart", pydeck_proto, layout_config=layout_config
            )

            return widget_state.value

        return self.dg._enqueue(
            "deck_gl_json_chart", pydeck_proto, layout_config=layout_config
        )

    @property
    def dg(self) -> DeltaGenerator:
        """The associated DeltaGenerator."""
        return cast("DeltaGenerator", self)


def _get_pydeck_width(pydeck_obj: Deck | None) -> int | None:
    """Extract the width from a pydeck Deck object, if specified."""
    if pydeck_obj is None:
        return None

    width = getattr(pydeck_obj, "width", None)
    if width is not None and isinstance(width, (int, float)):
        return int(width)

    return None


def _get_pydeck_tooltip(pydeck_obj: Deck | None) -> dict[str, str] | None:
    if pydeck_obj is None:
        return None

    # For pydeck <0.8.1 or pydeck>=0.8.1 when jupyter extra is installed.
    desk_widget = getattr(pydeck_obj, "deck_widget", None)
    if desk_widget is not None and isinstance(desk_widget.tooltip, dict):
        return desk_widget.tooltip

    # For pydeck >=0.8.1 when jupyter extra is not installed.
    # For details, see: https://github.com/visgl/deck.gl/pull/7125/files
    tooltip = getattr(pydeck_obj, "_tooltip", None)
    if tooltip is not None and isinstance(tooltip, dict):
        return cast("dict[str, str]", tooltip)

    return None


def _prepare_pydeck_for_json(pydeck_obj: Deck | None) -> None:
    """Prepare a pydeck Deck object for JSON serialization.

    This function converts pandas DataFrames in pydeck layers to lists of dicts
    to work around a pandas 3.x compatibility issue in pydeck's serialization.
    In pandas 3.x, DataFrames no longer have a __dict__ attribute that vars()
    can access, which breaks pydeck's default_serialize function.

    This function modifies the pydeck object in place. If the same Deck object
    is passed to multiple st.pydeck_chart calls within a single script run,
    subsequent calls will see the converted list[dict] data instead of DataFrames.
    In Streamlit's rerun-based execution model, this is typically not an issue
    since Deck objects are usually recreated on each run.

    For the upstream pydeck issue, see: https://github.com/visgl/deck.gl/issues/9986
    """
    if pydeck_obj is None:
        return

    import pandas as pd

    layers = getattr(pydeck_obj, "layers", None)
    if layers is None:
        return

    for layer in layers:
        data = getattr(layer, "data", None)
        if data is None:
            continue

        # Handle weakref to DataFrame (pydeck wraps DataFrames in weakrefs)
        if isinstance(data, weakref.ref):
            data = data()
            if data is None:
                continue

        # Convert pandas DataFrame to list of dicts for JSON serialization
        if isinstance(data, pd.DataFrame):
            layer.data = data.to_dict(orient="records")