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

from __future__ import annotations

from dataclasses import replace
from enum import Enum, EnumMeta
from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, cast, overload

from streamlit import config, logger
from streamlit.dataframe_util import OptionSequence, convert_anything_to_list
from streamlit.errors import StreamlitAPIException, StreamlitValueError
from streamlit.proto.SelectWidgetFilterMode_pb2 import (
    SelectWidgetFilterMode as ProtoSelectWidgetFilterMode,
)
from streamlit.runtime.state import get_session_state
from streamlit.type_util import (
    check_python_comparable,
)

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable, Sequence

    from streamlit.runtime.state.common import RegisterWidgetResult

_LOGGER: Final = logger.get_logger(__name__)

_FLOAT_EQUALITY_EPSILON: Final[float] = 0.000000000005
_Value = TypeVar("_Value")
T = TypeVar("T")

SelectWidgetFilterMode = Literal["fuzzy", "contains", "prefix"] | None

_VALID_SELECT_WIDGET_FILTER_MODES: Final = frozenset(
    {"fuzzy", "contains", "prefix", None}
)
_SELECT_WIDGET_FILTER_MODE_PROTO_MAP: Final = {
    "fuzzy": ProtoSelectWidgetFilterMode.FILTER_MODE_FUZZY,
    "contains": ProtoSelectWidgetFilterMode.FILTER_MODE_CONTAINS,
    "prefix": ProtoSelectWidgetFilterMode.FILTER_MODE_PREFIX,
    None: ProtoSelectWidgetFilterMode.FILTER_MODE_NONE,
}


def validate_select_widget_filter_mode(
    filter_mode: SelectWidgetFilterMode,
    *,
    accept_new_options: bool,
    command: Literal["st.selectbox", "st.multiselect"],
) -> ProtoSelectWidgetFilterMode.ValueType:
    """Validate ``filter_mode`` and return the protobuf enum value."""
    try:
        is_valid_filter_mode = filter_mode in _VALID_SELECT_WIDGET_FILTER_MODES
    except TypeError:
        is_valid_filter_mode = False

    if not is_valid_filter_mode:
        raise StreamlitValueError(
            "filter_mode",
            ["fuzzy", "contains", "prefix", "None"],
        )

    if filter_mode is None and accept_new_options:
        raise StreamlitAPIException(
            f"The `filter_mode` argument to `{command}` cannot be None when "
            "`accept_new_options=True`."
        )

    return _SELECT_WIDGET_FILTER_MODE_PROTO_MAP[filter_mode]


def index_(iterable: Iterable[_Value], x: _Value) -> int:
    """Return zero-based index of the first item whose value is equal to x.
    Raises a ValueError if there is no such item.

    We need a custom implementation instead of the built-in list .index() to
    be compatible with NumPy array and Pandas Series.

    Parameters
    ----------
    iterable : list, tuple, numpy.ndarray, pandas.Series
    x : Any

    Returns
    -------
    int
    """
    for i, value in enumerate(iterable):
        if x == value:
            return i
        if (
            isinstance(value, float)
            and isinstance(x, float)
            and abs(x - value) < _FLOAT_EQUALITY_EPSILON
        ):
            return i
    raise ValueError(f"{x} is not in iterable")


def check_and_convert_to_indices(
    opt: Sequence[Any], default_values: Sequence[Any] | Any | None
) -> list[int] | None:
    """Perform validation checks and return indices based on the default values."""
    if default_values is None:
        return None

    default_values = convert_anything_to_list(default_values)

    for value in default_values:
        if value not in opt:
            raise StreamlitAPIException(
                f"The default value '{value}' is not part of the options. "
                "Please make sure that every default values also exists in the options."
            )

    return [opt.index(value) for value in default_values]


def convert_to_sequence_and_check_comparable(options: OptionSequence[T]) -> Sequence[T]:
    indexable_options = convert_anything_to_list(options)
    check_python_comparable(indexable_options)
    return indexable_options


def get_default_indices(
    indexable_options: Sequence[T], default: Sequence[Any] | Any | None = None
) -> list[int]:
    default_indices = check_and_convert_to_indices(indexable_options, default)
    return default_indices if default_indices is not None else []


E1 = TypeVar("E1", bound=Enum)
E2 = TypeVar("E2", bound=Enum)

_ALLOWED_ENUM_COERCION_CONFIG_SETTINGS = ("off", "nameOnly", "nameAndValue")


def _coerce_enum(from_enum_value: E1, to_enum_class: type[E2]) -> E1 | E2:
    """Attempt to coerce an Enum value to another EnumMeta.

    An Enum value of EnumMeta E1 is considered coercible to EnumType E2
    if the EnumMeta __qualname__ match and the names of their members
    match as well. (This is configurable in streamlist configs)
    """
    if not isinstance(from_enum_value, Enum):
        raise ValueError(  # noqa: TRY004
            f"Expected an Enum in the first argument. Got {type(from_enum_value)}"
        )
    if not isinstance(to_enum_class, EnumMeta):
        raise ValueError(  # noqa: TRY004
            f"Expected an EnumMeta/Type in the second argument. Got {type(to_enum_class)}"
        )
    if isinstance(from_enum_value, to_enum_class):
        return from_enum_value  # Enum is already a member, no coercion necessary

    coercion_type = config.get_option("runner.enumCoercion")
    if coercion_type not in _ALLOWED_ENUM_COERCION_CONFIG_SETTINGS:
        raise StreamlitValueError(
            "runner.enumCoercion",
            [f"'{v}'" for v in _ALLOWED_ENUM_COERCION_CONFIG_SETTINGS],
        )
    if coercion_type == "off":
        return from_enum_value  # do not attempt to coerce

    # We now know this is an Enum AND the user has configured coercion enabled.
    # Check if we do NOT meet the required conditions and log a failure message
    # if that is the case.
    from_enum_class = from_enum_value.__class__
    if (
        from_enum_class.__qualname__ != to_enum_class.__qualname__
        or (
            coercion_type == "nameOnly"
            and set(to_enum_class._member_names_) != set(from_enum_class._member_names_)
        )
        or (
            coercion_type == "nameAndValue"
            and set(to_enum_class._value2member_map_)
            != set(from_enum_class._value2member_map_)
        )
    ):
        _LOGGER.debug("Failed to coerce %s to class %s", from_enum_value, to_enum_class)
        return from_enum_value  # do not attempt to coerce

    # At this point we think the Enum is coercible, and we know
    # E1 and E2 have the same member names. We convert from E1 to E2 using _name_
    # (since user Enum subclasses can override the .name property in 3.11)
    _LOGGER.debug("Coerced %s to class %s", from_enum_value, to_enum_class)
    return to_enum_class[from_enum_value._name_]


def _extract_common_class_from_iter(iterable: Iterable[Any]) -> Any:
    """Return the common class of all elements in a iterable if they share one.
    Otherwise, return None.
    """
    try:
        inner_iter = iter(iterable)
        first_class = type(next(inner_iter))
    except StopIteration:
        return None
    if all(type(item) is first_class for item in inner_iter):
        return first_class
    return None


@overload
def maybe_coerce_enum(
    register_widget_result: RegisterWidgetResult[Enum],
    options: type[Enum],
    opt_sequence: Sequence[Any],
) -> RegisterWidgetResult[Enum]: ...


@overload
def maybe_coerce_enum(
    register_widget_result: RegisterWidgetResult[T],
    options: OptionSequence[T],
    opt_sequence: Sequence[T],
) -> RegisterWidgetResult[T]: ...


def maybe_coerce_enum(
    register_widget_result: RegisterWidgetResult[Any],
    options: OptionSequence[Any],
    opt_sequence: Sequence[Any],
) -> RegisterWidgetResult[Any]:
    """Maybe Coerce a RegisterWidgetResult with an Enum member value to
    RegisterWidgetResult[option] if option is an EnumType, otherwise just return
    the original RegisterWidgetResult.
    """

    value = register_widget_result.value

    # If the value is not an Enum, return early
    if not isinstance(value, Enum):
        return register_widget_result

    coerce_class: EnumMeta | None
    if isinstance(options, EnumMeta):
        coerce_class = options
    else:
        coerce_class = _extract_common_class_from_iter(opt_sequence)
        if coerce_class is None:
            return register_widget_result
    enum_class = cast("type[Enum]", coerce_class)

    # Use replace so other fields (e.g. incoming_serialized_value) are preserved
    # rather than dropped when only the value is coerced.
    return replace(
        register_widget_result,
        value=_coerce_enum(value, enum_class),
    )


# slightly ugly typing because TypeVars with Generic Bounds are not supported
# (https://github.com/python/typing/issues/548)
@overload
def maybe_coerce_enum_sequence(
    register_widget_result: RegisterWidgetResult[list[T] | list[T | str]],
    options: OptionSequence[T],
    opt_sequence: Sequence[T],
) -> RegisterWidgetResult[list[T] | list[T | str]]: ...


@overload
def maybe_coerce_enum_sequence(
    register_widget_result: RegisterWidgetResult[tuple[T, T]],
    options: OptionSequence[T],
    opt_sequence: Sequence[T],
) -> RegisterWidgetResult[tuple[T, T]]: ...


def maybe_coerce_enum_sequence(
    register_widget_result: RegisterWidgetResult[list[Any] | tuple[Any, ...]],
    options: OptionSequence[Any],
    opt_sequence: Sequence[Any],
) -> RegisterWidgetResult[list[Any] | tuple[Any, ...]]:
    """Maybe Coerce a RegisterWidgetResult with a sequence of Enum members as value
    to RegisterWidgetResult[Sequence[option]] if option is an EnumType, otherwise just
    return the original RegisterWidgetResult.
    """

    # If not all widget values are Enums, return early
    if not all(isinstance(val, Enum) for val in register_widget_result.value):
        return register_widget_result

    # Extract the class to coerce
    coerce_class: EnumMeta | None
    if isinstance(options, EnumMeta):
        coerce_class = options
    else:
        coerce_class = _extract_common_class_from_iter(opt_sequence)
        if coerce_class is None:
            return register_widget_result
    enum_class = cast("type[Enum]", coerce_class)

    # Return a new RegisterWidgetResult with the coerced enum values sequence.
    # Use replace so other fields (e.g. incoming_serialized_value) are preserved.
    return replace(
        register_widget_result,
        value=type(register_widget_result.value)(
            _coerce_enum(cast("Enum", val), enum_class)
            for val in register_widget_result.value
        ),
    )


def create_mappings(
    options: Sequence[T], format_func: Callable[[T], str] = str
) -> tuple[list[str], dict[str, int]]:
    """Iterates through the options and formats them using the format_func.

    Returns a tuple of the formatted options and a mapping of the formatted options to
    the original options.
    """
    formatted_option_to_option_mapping: dict[str, int] = {}
    formatted_options: list[str] = []
    for index, option in enumerate(options):
        formatted_option = format_func(option)
        formatted_options.append(formatted_option)
        # If formatted labels are duplicated, the last one wins. We keep this
        # behavior to mirror radio/selectbox/multiselect, but it makes selection
        # ambiguous for string-based widgets.
        # TODO: Consider raising a StreamlitAPIException on duplicate labels.
        formatted_option_to_option_mapping[formatted_option] = index

    return (
        formatted_options,
        formatted_option_to_option_mapping,
    )


def validate_and_sync_value_with_options(
    current_value: T | None,
    opt: Sequence[T],
    default_index: int | None,
    key: str | int | None,
    format_func: Callable[[Any], str] = str,
) -> tuple[T | None, bool]:
    """Validate current value against options, resetting session state if invalid.

    This function has a side-effect: if the value is not found in the options
    and a key is provided, it will update session state with the new value.


    Parameters
    ----------
    current_value
        The current widget value to validate.
    opt
        The sequence of valid options.
    default_index
        The default index to reset to if value is invalid.
    key
        The widget key for session state updates.
    format_func
        Function to format options for comparison. Used to compare values by their
        string representation instead of using == directly. This is necessary because
        widget values are deepcopied, and for custom classes without __eq__, the
        deepcopied instances would fail identity comparison.

    Returns
    -------
    tuple[T | None, bool]
        A tuple of (validated_value, value_was_reset).
    """
    if current_value is None:
        return current_value, False

    # Use format_func comparison for all values. This correctly handles:
    # - Custom objects without __eq__ (deepcopied instances)
    # - Enum values (already from current class due to serde deserialization)
    formatted_options_set = {format_func(o) for o in opt}
    try:
        formatted_value = format_func(current_value)
        if formatted_value in formatted_options_set:
            return current_value, False
    except Exception:  # noqa: S110
        pass  # format_func failed - value is invalid, fall through to reset

    # Value not in options - reset to default
    if default_index is not None and len(opt) > 0:
        new_value: T | None = opt[default_index]
    else:
        new_value = None

    if key is not None:
        # Update session_state so subsequent accesses in this run
        # return the corrected value. Use reset_state_value to avoid
        # the "cannot be modified after widget instantiated" error.
        get_session_state().reset_state_value(str(key), new_value)

    return new_value, True


def resolve_value_against_options(
    current_value: T | None,
    opt: Sequence[T],
    formatted_option_to_option_index: dict[str, int],
    default_index: int | None,
    key: str | int | None,
    format_func: Callable[[Any], str] = str,
    incoming_serialized_value: str | None = None,
) -> tuple[T | None, bool]:
    """Like :func:`validate_and_sync_value_with_options`, but falls back to the
    stored wire label when ``format_func`` raises on the current value.

    Membership is decided primarily by ``format_func``: the value is valid if
    its formatted label is among the options. This preserves correct behavior
    when ``format_func`` changes between runs or when the value was set
    programmatically.

    If ``format_func`` raises on ``current_value`` -- which happens when it
    performs an identity- or class-dependent operation (e.g. a dict lookup) and
    the value is a deepcopy stored from an earlier run -- the stored wire label
    ``incoming_serialized_value`` is used instead to decide membership. On a
    match, the matching *current* option instance is returned so the value
    always belongs to the current run rather than a stale deepcopy.

    Has the same session-state side effect as
    :func:`validate_and_sync_value_with_options`: when the value is reset and a
    key is provided, session state is updated with the new value.

    Parameters
    ----------
    current_value
        The current widget value to validate.
    opt
        The sequence of valid options.
    formatted_option_to_option_index
        Mapping from formatted option labels to their index in ``opt``.
    default_index
        The default index to reset to if the value is invalid.
    key
        The widget key for session state updates.
    format_func
        Used to compute the value's label for membership. May raise on a stale
        value, which triggers the wire-label fallback.
    incoming_serialized_value
        The widget's stored wire label, used as a fallback identity when
        ``format_func`` raises on ``current_value``. ``None`` if unavailable.

    Returns
    -------
    tuple[T | None, bool]
        A tuple of (validated_value, value_was_reset).
    """
    if current_value is None:
        return current_value, False

    try:
        formatted_value = format_func(current_value)
        if formatted_value in formatted_option_to_option_index:
            return current_value, False
    except Exception:
        # format_func can't handle the stored value (e.g. it's identity- or
        # class-dependent and the value is a deepcopy from an earlier run). Use
        # the stored wire label as the identity instead, and return the current
        # option instance so we never hand back a stale deepcopy.
        option_index = (
            formatted_option_to_option_index.get(incoming_serialized_value)
            if incoming_serialized_value is not None
            else None
        )
        if option_index is not None and 0 <= option_index < len(opt):
            return opt[option_index], False

    # Value not in options - reset to default.
    if default_index is not None and len(opt) > 0:
        new_value: T | None = opt[default_index]
    else:
        new_value = None

    if key is not None:
        get_session_state().reset_state_value(str(key), new_value)

    return new_value, True


def validate_and_sync_multiselect_value_with_options(
    current_values: list[T] | list[T | str],
    opt: Sequence[T],
    key: str | int | None,
    format_func: Callable[[Any], str] = str,
) -> tuple[list[T] | list[T | str], bool]:
    """Validate multiselect values against options, syncing session state if needed.

    This function has a side-effect: if any values are filtered out and a key
    is provided, it will update session state with the filtered list.

    Unlike selectbox which resets to a default when the value is invalid,
    multiselect filters out invalid values and keeps the valid ones.

    Parameters
    ----------
    current_values
        The current list of selected values to validate.
    opt
        The sequence of valid options.
    key
        The widget key for session state updates.
    format_func
        Function to format options for comparison. Used to compare values by their
        string representation instead of using == directly. This is necessary because
        widget values are deepcopied, and for custom classes without __eq__, the
        deepcopied instances would fail identity comparison.

    Returns
    -------
    tuple[list[T] | list[T | str], bool]
        A tuple of (validated_values, values_were_filtered).
    """
    if not current_values:
        return current_values, False

    # Create a set of formatted options for O(1) lookup.
    # We use format_func to compare values by their string representation
    # instead of using == directly. This is necessary because widget values
    # are deepcopied, and for custom classes without __eq__, the deepcopied
    # instances would fail identity comparison.
    formatted_options_set = {format_func(o) for o in opt}

    valid_values: list[T | str] = []
    for value in current_values:
        try:
            formatted_value = format_func(value)
        except Exception:  # noqa: S112
            # format_func failed on this value (e.g., a string value from a previous
            # session when format_func expects an object with specific attributes).
            # In this case, the value is definitely not valid since the current options
            # can be formatted successfully.
            continue

        if formatted_value in formatted_options_set:
            valid_values.append(value)

    if len(valid_values) == len(current_values):
        return current_values, False

    if key is not None:
        get_session_state().reset_state_value(str(key), valid_values)

    return valid_values, True


def validate_and_sync_range_value_with_options(
    current_value: tuple[T, T],
    opt: Sequence[T],
    default_indices: list[int],
    key: str | int | None,
    format_func: Callable[[Any], str] = str,
) -> tuple[tuple[T, T], bool]:
    """Validate a range value (tuple of two values) against options.

    If either value in the range is not found in options, the entire range is
    reset to the default. This function has a side-effect: if the values are
    invalid and a key is provided, it will update session state with the new value.

    Parameters
    ----------
    current_value
        The current range value (tuple of two values) to validate.
    opt
        The sequence of valid options.
    default_indices
        The default indices to reset to if value is invalid. Should contain
        at least one index; if only one index is provided, the second default
        will be the last option.
    key
        The widget key for session state updates.
    format_func
        Function to format options for comparison. Used to compare values by their
        string representation instead of using == directly.

    Returns
    -------
    tuple[tuple[T, T], bool]
        A tuple of (validated_value, value_was_reset).
    """
    if len(opt) == 0:
        return current_value, False

    formatted_options_set = {format_func(o) for o in opt}

    def is_valid(val: Any) -> bool:
        """Check if a value exists in options via format_func comparison."""
        try:
            return format_func(val) in formatted_options_set
        except Exception:
            return False

    def get_default_range() -> tuple[T, T]:
        """Get the default range value."""
        end_idx = default_indices[1] if len(default_indices) > 1 else len(opt) - 1
        return (opt[default_indices[0]], opt[end_idx])

    # Validate both values in the range.
    if is_valid(current_value[0]) and is_valid(current_value[1]):
        return current_value, False

    # Either value is invalid - reset entire range.
    new_value = get_default_range()
    if key is not None:
        get_session_state().reset_state_value(str(key), new_value)
    return new_value, True
