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

"""Runtime-related utility functions."""

from __future__ import annotations

from logging import getLogger
from typing import TYPE_CHECKING, Any, Final, cast

from streamlit import config
from streamlit.errors import MarkdownFormattedException

if TYPE_CHECKING:
    from streamlit.proto.ForwardMsg_pb2 import ForwardMsg

_LOGGER: Final = getLogger(__name__)

# The interval (in seconds) between message flush cycles in the runtime event loop.
# Also reused by callers that need to wait for outstanding messages to be delivered
# before requesting a rerun (e.g. `st.switch_page`).
#
# Message coalescing is implemented at the queue level (`_maybe_compose_delta_msgs`),
# but this interval is the window during which it can happen. A small non-zero value
# is important for single-element placeholders like `st.empty` and `st.skeleton`:
# when a placeholder is shown and then replaced within the same run, both deltas
# share a delta_path and coalesce into the final element, so the browser never
# renders the intermediate state. A too-small interval can flush the placeholder
# first and briefly unmount/remount the real element (e.g. flash a skeleton) on
# reruns.
MESSAGE_FLUSH_INTERVAL_SECS: Final[float] = 0.005


class MessageSizeError(MarkdownFormattedException):
    """Exception raised when a websocket message is larger than the configured limit."""

    def __init__(self, failed_msg_str: Any) -> None:
        msg = self._get_message(failed_msg_str)
        super().__init__(msg)

    def _get_message(self, failed_msg_str: Any) -> str:
        # This needs to have zero indentation otherwise the markdown will render incorrectly.
        return (
            f"""
**Data of size {len(failed_msg_str) / 1e6:.1f} MB exceeds the message size limit of
{get_max_message_size_bytes() / 1e6} MB.**

This is often caused by a large chart or dataframe. Please decrease the amount of data sent
to the browser, or increase the limit by setting the config option `server.maxMessageSize`.
[Click here to learn more about config options](https://docs.streamlit.io/develop/api-reference/configuration/config.toml).

_Note that increasing the limit may lead to long loading times and large memory consumption
of the client's browser and the Streamlit server._
"""
        ).strip("\n")


class WidgetStateSizeError(MarkdownFormattedException):
    """Exception raised when client-sent widget state is larger than the limit."""

    def __init__(self, failed_msg_size: int) -> None:
        msg = self._get_message(failed_msg_size)
        super().__init__(msg)

    def _get_message(self, failed_msg_size: int) -> str:
        # This needs to have zero indentation otherwise the markdown will render incorrectly.
        return (
            f"""
**Widget state of size {failed_msg_size / 1e6:.1f} MB exceeds the widget state size limit of
{get_max_widget_state_size_bytes() / 1e6} MB.**

This can happen when a widget or custom component sends a large value to Streamlit.
Please decrease the amount of data sent from the browser, or increase the limit by
setting the config option `server.maxWidgetStateSize`.
[Click here to learn more about config options](https://docs.streamlit.io/develop/api-reference/configuration/config.toml).

_Note that increasing the limit may lead to long loading times and large memory consumption
of the Streamlit server._
"""
        ).strip("\n")


def serialize_forward_msg(msg: ForwardMsg) -> bytes:
    """Serialize a ForwardMsg to send to a client.

    If the message is too large, it will be converted to an exception message
    instead.
    """

    msg_str = msg.SerializeToString()

    if len(msg_str) > get_max_message_size_bytes():
        # Overwrite the offending ForwardMsg.delta with an error to display.
        # This assumes that the size limit wasn't exceeded due to metadata.
        from streamlit.elements import exception

        msg_size_error = MessageSizeError(msg_str)
        _LOGGER.warning(
            "Websocket message size limit exceeded. Showing error to the user: %s",
            msg_size_error,
        )
        exception.marshall(msg.delta.new_element.exception, msg_size_error)
        # Deactivate caching for this error message:
        msg.metadata.cacheable = False
        msg_str = msg.SerializeToString()

    return msg_str


# This needs to be initialized lazily to avoid calling config.get_option() and
# thus initializing config options when this file is first imported.
_max_message_size_bytes: int | None = None
_max_widget_state_size_bytes: int | None = None


def get_max_message_size_bytes() -> int:
    """Returns the max websocket message size in bytes.

    This will lazyload the value from the config and store it in the global symbol table.
    """
    global _max_message_size_bytes  # noqa: PLW0603

    if _max_message_size_bytes is None:
        _max_message_size_bytes = config.get_option("server.maxMessageSize") * int(1e6)

    return cast("int", _max_message_size_bytes)


def get_max_widget_state_size_bytes() -> int:
    """Returns the max client-sent message size in bytes.

    This limit is applied both to raw inbound WebSocket frames in the Starlette
    handler and to the aggregate widget state protobuf in the runtime layer.

    This will lazyload the value from the config and store it in the global symbol table.
    """
    global _max_widget_state_size_bytes  # noqa: PLW0603

    if _max_widget_state_size_bytes is None:
        _max_widget_state_size_bytes = config.get_option(
            "server.maxWidgetStateSize"
        ) * int(1e6)

    return cast("int", _max_widget_state_size_bytes)