# 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

import copy
import json
from collections.abc import Mapping
from dataclasses import dataclass
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    Literal,
    TypeAlias,
    overload,
)

from streamlit.dataframe_util import DataFormat
from streamlit.elements.lib.column_types import (
    ButtonColumnResult,
    ColumnConfig,
    ColumnType,
)
from streamlit.elements.lib.dicttools import remove_none_values
from streamlit.elements.lib.policies import check_widget_policies
from streamlit.elements.lib.utils import compute_and_register_element_id
from streamlit.errors import StreamlitAPIException
from streamlit.runtime.state import register_widget
from streamlit.util import ReadOnlyAttributeDictionary

if TYPE_CHECKING:
    import pyarrow as pa
    from pandas import DataFrame, Index, Series

    from streamlit.delta_generator import DeltaGenerator
    from streamlit.proto.Common_pb2 import StringTriggerValue
    from streamlit.proto.Dataframe_pb2 import Dataframe as DataframeProto
    from streamlit.runtime.scriptrunner_utils.script_run_context import ScriptRunContext


# The index identifier can be used to apply configuration options
IndexIdentifierType = Literal["_index"]
INDEX_IDENTIFIER: IndexIdentifierType = "_index"

# This is used as prefix for columns that are configured via the numerical position.
# The integer value is converted into a string key with this prefix.
# This needs to match with the prefix configured in the frontend.
NUMERICAL_POSITION_PREFIX = "_pos:"


# The column data kind is used to describe the type of the data within the column.
class ColumnDataKind(str, Enum):
    INTEGER = "integer"
    FLOAT = "float"
    DATE = "date"
    TIME = "time"
    DATETIME = "datetime"
    BOOLEAN = "boolean"
    STRING = "string"
    TIMEDELTA = "timedelta"
    PERIOD = "period"
    INTERVAL = "interval"
    BYTES = "bytes"
    DECIMAL = "decimal"
    COMPLEX = "complex"
    LIST = "list"
    DICT = "dict"
    EMPTY = "empty"
    UNKNOWN = "unknown"


# The dataframe schema is a mapping from the name of the column
# in the underlying dataframe to the column data kind.
# The index column uses `_index` as name.
DataframeSchema: TypeAlias = dict[str, ColumnDataKind]

# This mapping contains all editable column types mapped to the data kinds
# that the column type is compatible for editing.
_EDITING_COMPATIBILITY_MAPPING: Final[dict[ColumnType, list[ColumnDataKind]]] = {
    "text": [ColumnDataKind.STRING, ColumnDataKind.EMPTY],
    "number": [
        ColumnDataKind.INTEGER,
        ColumnDataKind.FLOAT,
        ColumnDataKind.DECIMAL,
        ColumnDataKind.STRING,
        ColumnDataKind.TIMEDELTA,
        ColumnDataKind.EMPTY,
    ],
    "checkbox": [
        ColumnDataKind.BOOLEAN,
        ColumnDataKind.STRING,
        ColumnDataKind.INTEGER,
        ColumnDataKind.EMPTY,
    ],
    "selectbox": [
        ColumnDataKind.STRING,
        ColumnDataKind.BOOLEAN,
        ColumnDataKind.INTEGER,
        ColumnDataKind.FLOAT,
        ColumnDataKind.EMPTY,
    ],
    "date": [ColumnDataKind.DATE, ColumnDataKind.DATETIME, ColumnDataKind.EMPTY],
    "time": [ColumnDataKind.TIME, ColumnDataKind.DATETIME, ColumnDataKind.EMPTY],
    "datetime": [
        ColumnDataKind.DATETIME,
        ColumnDataKind.DATE,
        ColumnDataKind.TIME,
        ColumnDataKind.EMPTY,
    ],
    "link": [ColumnDataKind.STRING, ColumnDataKind.EMPTY],
    "list": [
        ColumnDataKind.LIST,
        ColumnDataKind.STRING,
        ColumnDataKind.EMPTY,
    ],
    "multiselect": [
        ColumnDataKind.LIST,
        ColumnDataKind.STRING,
        ColumnDataKind.EMPTY,
    ],
    "markdown": [ColumnDataKind.STRING, ColumnDataKind.EMPTY],
}


def is_type_compatible(column_type: ColumnType, data_kind: ColumnDataKind) -> bool:
    """Check if the column type is compatible with the underlying data kind.

    This check only applies to editable column types (e.g. number or text).
    Non-editable column types (e.g. bar_chart or image) can be configured for
    all data kinds (this might change in the future).

    Parameters
    ----------
    column_type : ColumnType
        The column type to check.

    data_kind : ColumnDataKind
        The data kind to check.

    Returns
    -------
    bool
        True if the column type is compatible with the data kind, False otherwise.
    """

    if column_type not in _EDITING_COMPATIBILITY_MAPPING:
        return True

    return data_kind in _EDITING_COMPATIBILITY_MAPPING[column_type]


def _determine_data_kind_via_arrow(field: pa.Field) -> ColumnDataKind:
    """Determine the data kind via the arrow type information.

    The column data kind refers to the shared data type of the values
    in the column (e.g. int, float, str, bool).

    Parameters
    ----------
    field : pa.Field
        The arrow field from the arrow table schema.

    Returns
    -------
    ColumnDataKind
        The data kind of the field.
    """
    import pyarrow as pa

    field_type = field.type
    if pa.types.is_integer(field_type):
        return ColumnDataKind.INTEGER

    if pa.types.is_floating(field_type):
        return ColumnDataKind.FLOAT

    if pa.types.is_boolean(field_type):
        return ColumnDataKind.BOOLEAN

    if pa.types.is_string(field_type) or pa.types.is_large_string(field_type):
        return ColumnDataKind.STRING

    if pa.types.is_date(field_type):
        return ColumnDataKind.DATE

    if pa.types.is_time(field_type):
        return ColumnDataKind.TIME

    if pa.types.is_timestamp(field_type):
        return ColumnDataKind.DATETIME

    if pa.types.is_duration(field_type):
        return ColumnDataKind.TIMEDELTA

    if pa.types.is_list(field_type):
        return ColumnDataKind.LIST

    if pa.types.is_decimal(field_type):
        return ColumnDataKind.DECIMAL

    if pa.types.is_null(field_type):
        return ColumnDataKind.EMPTY

    # Interval does not seem to work correctly:
    # if pa.types.is_interval(field_type):
    #     return ColumnDataKind.INTERVAL  # noqa: ERA001

    if pa.types.is_binary(field_type):
        return ColumnDataKind.BYTES

    if pa.types.is_struct(field_type):
        return ColumnDataKind.DICT

    return ColumnDataKind.UNKNOWN


def _determine_data_kind_via_pandas_dtype(
    column: Series[Any] | Index[Any],
) -> ColumnDataKind:
    """Determine the data kind by using the pandas dtype.

    The column data kind refers to the shared data type of the values
    in the column (e.g. int, float, str, bool).

    Parameters
    ----------
    column : pd.Series, pd.Index
        The column for which the data kind should be determined.

    Returns
    -------
    ColumnDataKind
        The data kind of the column.
    """
    import pandas as pd

    column_dtype = column.dtype
    if pd.api.types.is_bool_dtype(column_dtype):
        return ColumnDataKind.BOOLEAN

    if pd.api.types.is_integer_dtype(column_dtype):
        return ColumnDataKind.INTEGER

    if pd.api.types.is_float_dtype(column_dtype):
        return ColumnDataKind.FLOAT

    if pd.api.types.is_datetime64_any_dtype(column_dtype):
        return ColumnDataKind.DATETIME

    if pd.api.types.is_timedelta64_dtype(column_dtype):
        return ColumnDataKind.TIMEDELTA

    if isinstance(column_dtype, pd.PeriodDtype):
        return ColumnDataKind.PERIOD

    if isinstance(column_dtype, pd.IntervalDtype):
        return ColumnDataKind.INTERVAL

    if pd.api.types.is_complex_dtype(column_dtype):
        return ColumnDataKind.COMPLEX

    if pd.api.types.is_object_dtype(
        column_dtype
    ) is False and pd.api.types.is_string_dtype(column_dtype):
        # This handles pandas 3.0+ StringDtype (and PyArrow-backed string types).
        # We exclude object dtype here because object columns with string values
        # are handled via _determine_data_kind_via_inferred_type in the caller.
        return ColumnDataKind.STRING

    return ColumnDataKind.UNKNOWN


def _determine_data_kind_via_inferred_type(
    column: Series[Any] | Index[Any],
) -> ColumnDataKind:
    """Determine the data kind by inferring it from the underlying data.

    The column data kind refers to the shared data type of the values
    in the column (e.g. int, float, str, bool).

    Parameters
    ----------
    column : pd.Series, pd.Index
        The column to determine the data kind for.

    Returns
    -------
    ColumnDataKind
        The data kind of the column.
    """
    from pandas.api.types import infer_dtype

    inferred_type = infer_dtype(column)

    if inferred_type == "string":
        return ColumnDataKind.STRING

    if inferred_type == "bytes":
        return ColumnDataKind.BYTES

    if inferred_type in {"floating", "mixed-integer-float"}:
        return ColumnDataKind.FLOAT

    if inferred_type == "integer":
        return ColumnDataKind.INTEGER

    if inferred_type == "decimal":
        return ColumnDataKind.DECIMAL

    if inferred_type == "complex":
        return ColumnDataKind.COMPLEX

    if inferred_type == "boolean":
        return ColumnDataKind.BOOLEAN

    if inferred_type in {"datetime64", "datetime"}:
        return ColumnDataKind.DATETIME

    if inferred_type == "date":
        return ColumnDataKind.DATE

    if inferred_type in {"timedelta64", "timedelta"}:
        return ColumnDataKind.TIMEDELTA

    if inferred_type == "time":
        return ColumnDataKind.TIME

    if inferred_type == "period":
        return ColumnDataKind.PERIOD

    if inferred_type == "interval":
        return ColumnDataKind.INTERVAL

    if inferred_type == "empty":
        return ColumnDataKind.EMPTY

    # Unused types: mixed, unknown-array, categorical, mixed-integer

    return ColumnDataKind.UNKNOWN


def _determine_data_kind(
    column: Series[Any] | Index[Any], field: pa.Field | None = None
) -> ColumnDataKind:
    """Determine the data kind of a column.

    The column data kind refers to the shared data type of the values
    in the column (e.g. int, float, str, bool).

    Parameters
    ----------
    column : pd.Series, pd.Index
        The column to determine the data kind for.
    field : pa.Field, optional
        The arrow field from the arrow table schema.

    Returns
    -------
    ColumnDataKind
        The data kind of the column.
    """
    import pandas as pd

    if isinstance(column.dtype, pd.CategoricalDtype):
        # Categorical columns can have different underlying data kinds
        # depending on the categories.
        return _determine_data_kind_via_inferred_type(column.dtype.categories)

    if field is not None:
        data_kind = _determine_data_kind_via_arrow(field)
        if data_kind != ColumnDataKind.UNKNOWN:
            return data_kind

    if column.dtype.name == "object":
        # If dtype is object, we need to infer the type from the column
        return _determine_data_kind_via_inferred_type(column)
    return _determine_data_kind_via_pandas_dtype(column)


def determine_dataframe_schema(
    data_df: DataFrame, arrow_schema: pa.Schema
) -> DataframeSchema:
    """Determine the schema of a dataframe.

    Parameters
    ----------
    data_df : pd.DataFrame
        The dataframe to determine the schema of.
    arrow_schema : pa.Schema
        The Arrow schema of the dataframe.

    Returns
    -------
    DataframeSchema
        A mapping that contains the detected data type for the index and columns.
        The key is the column name in the underlying dataframe or ``_index`` for index columns.
    """

    dataframe_schema: DataframeSchema = {}

    # Add type of index:
    # TODO(lukasmasuch): We need to apply changes here to support multiindex.
    dataframe_schema[INDEX_IDENTIFIER] = _determine_data_kind(data_df.index)

    # Add types for all columns:
    for i, column in enumerate(data_df.items()):
        column_name = str(column[0])
        column_data = column[1]
        dataframe_schema[column_name] = _determine_data_kind(
            column_data, arrow_schema.field(i)
        )
    return dataframe_schema


# A mapping of column names/IDs to column configs.
ColumnConfigMapping: TypeAlias = dict[IndexIdentifierType | str | int, ColumnConfig]

ColumnConfigMappingInput: TypeAlias = Mapping[
    # TODO(lukasmasuch): This should also use int here to
    # correctly type the support for positional index. However,
    # allowing int here leads mypy to complain about simple dict[str, ...]
    # as input -> which seems like a mypy bug.
    IndexIdentifierType | str,
    ColumnConfig | ButtonColumnResult | str | None,
]


ButtonColumnMapping: TypeAlias = dict[str, ButtonColumnResult]


class ButtonColumnClickState(ReadOnlyAttributeDictionary):
    """The schema for button click state in ButtonColumn.

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

    Read-only dict-like click payload with attribute and key access
    (``click.row`` / ``click["row"]``). Both fields are always present when a
    button click occurs.
    """

    row: int
    label: str

    @overload
    def __getitem__(self, key: Literal["row"]) -> int: ...

    @overload
    def __getitem__(self, key: Literal["label"]) -> str: ...

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

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


@dataclass
class ButtonClickSerde:
    """Serializer/deserializer for ButtonColumn click values.

    Uses string trigger value pattern (value resets after each run).
    The frontend sends the click state as a JSON string.
    """

    def serialize(self, v: ButtonColumnClickState | None) -> StringTriggerValue:
        from streamlit.proto.Common_pb2 import StringTriggerValue

        if v is None:
            return StringTriggerValue()
        return StringTriggerValue(data=json.dumps(v))

    def deserialize(self, ui_value: str | None) -> ButtonColumnClickState | None:
        if not ui_value:
            return None

        parsed = json.loads(ui_value)
        # Validate shape: must be a dict with "row" (int >= 0) and "label" (str).
        # bool is a subclass of int, so booleans are explicitly rejected for "row".
        if (
            not isinstance(parsed, dict)
            or not isinstance(parsed.get("row"), int)
            or isinstance(parsed.get("row"), bool)
            or not isinstance(parsed.get("label"), str)
        ):
            raise StreamlitAPIException(
                "Invalid button click state: expected {row: int, label: str}."
            )

        # Validate row is non-negative (bounds check - row < num_rows is
        # checked downstream when accessing the data)
        if parsed["row"] < 0:
            raise StreamlitAPIException(
                f"Invalid button click row index: {parsed['row']}. Row must be >= 0."
            )

        return ButtonColumnClickState(parsed)


def extract_button_column_configs(
    column_config: ColumnConfigMappingInput | None,
) -> tuple[dict[Any, ColumnConfig | str | None] | None, ButtonColumnMapping]:
    """Extract interactive button-column configs from a column config mapping."""
    button_columns: ButtonColumnMapping = {}
    processed_column_config: dict[Any, ColumnConfig | str | None] | None = None

    if column_config is not None:
        processed_column_config = {}
        for col_name, config in column_config.items():
            if isinstance(config, ButtonColumnResult):
                # Transform key the same way column config does for consistency.
                column_widget_key = (
                    f"{NUMERICAL_POSITION_PREFIX}{col_name}"
                    if isinstance(col_name, int)
                    else str(col_name)
                )
                button_columns[column_widget_key] = config
                processed_column_config[col_name] = config.config
            else:
                processed_column_config[col_name] = config

    return processed_column_config, button_columns


def register_button_column_widgets(
    *,
    dg: DeltaGenerator,
    proto: DataframeProto,
    button_columns: ButtonColumnMapping,
    ctx: ScriptRunContext | None,
) -> None:
    """Register widgets for interactive button columns and attach them to a dataframe proto."""
    button_serde = ButtonClickSerde()
    for col_name, button_col in button_columns.items():
        check_widget_policies(
            dg,
            button_col.key,
            on_change=button_col.on_click,
            default_value=None,
            writes_allowed=False,
        )
        widget_id = compute_and_register_element_id(
            "dataframe_button",
            user_key=button_col.key,
            key_as_main_identity=True,
            dg=dg,
        )
        register_widget(
            widget_id,
            on_change_handler=button_col.on_click,
            args=button_col.args,
            kwargs=button_col.kwargs,
            deserializer=button_serde.deserialize,
            serializer=button_serde.serialize,
            ctx=ctx,
            value_type="string_trigger_value",
        )
        proto.button_click_widgets[col_name] = widget_id


def process_config_mapping(
    column_config: ColumnConfigMappingInput | None = None,
) -> ColumnConfigMapping:
    """Transforms a user-provided column config mapping into a valid column config mapping
    that can be used by the frontend.

    Parameters
    ----------
    column_config: dict or None
        The user-provided column config mapping.

    Returns
    -------
    dict
        The transformed column config mapping.
    """
    if column_config is None:
        return {}

    transformed_column_config: ColumnConfigMapping = {}
    for column, config in column_config.items():
        if config is None:
            transformed_column_config[column] = ColumnConfig(hidden=True)
        elif isinstance(config, str):
            transformed_column_config[column] = ColumnConfig(label=config)
        elif isinstance(config, ButtonColumnResult):
            # ButtonColumnResult is typically preprocessed before reaching this function
            # in both st.dataframe and st.data_editor. If we encounter it here, extract
            # the config dict. Button columns are always read-only.
            transformed_column_config[column] = copy.deepcopy(config.config)
        elif isinstance(config, dict):
            # Ensure that the column config objects are cloned
            # since we will apply in-place changes to it.
            transformed_column_config[column] = copy.deepcopy(config)
        else:
            raise StreamlitAPIException(
                f"Invalid column config for column `{column}`. "
                f"Expected `None`, `str` or `dict`, but got `{type(config)}`."
            )
    return transformed_column_config


def update_column_config(
    column_config_mapping: ColumnConfigMapping,
    column: str | int,
    column_config: ColumnConfig,
) -> None:
    """Updates the column config value for a single column within the mapping.

    Parameters
    ----------
    column_config_mapping : ColumnConfigMapping
        The column config mapping to update.

    column : str | int
        The column to update the config value for. This can be the column name or
        the numerical position of the column.

    column_config : ColumnConfig
        The column config to update.
    """

    if column not in column_config_mapping:
        column_config_mapping[column] = {}

    column_config_mapping[column].update(column_config)


def apply_data_specific_configs(
    columns_config: ColumnConfigMapping,
    data_format: DataFormat,
) -> None:
    """Apply data specific configurations to the provided dataframe.

    This will apply inplace changes to the dataframe and the column configurations
    depending on the data format.

    Parameters
    ----------
    columns_config : ColumnConfigMapping
        A mapping of column names/ids to column configurations.

    data_format : DataFormat
        The format of the data.
    """

    # Pandas adds a range index as default to all datastructures
    # but for most of the non-pandas data objects it is unnecessary
    # to show this index to the user. Therefore, we will hide it as default.
    if data_format in {
        DataFormat.SET_OF_VALUES,
        DataFormat.TUPLE_OF_VALUES,
        DataFormat.LIST_OF_VALUES,
        DataFormat.NUMPY_LIST,
        DataFormat.NUMPY_MATRIX,
        DataFormat.LIST_OF_RECORDS,
        DataFormat.LIST_OF_ROWS,
        DataFormat.COLUMN_VALUE_MAPPING,
        # Dataframe-like objects that don't have an index:
        DataFormat.PANDAS_ARRAY,
        DataFormat.PANDAS_INDEX,
        DataFormat.POLARS_DATAFRAME,
        DataFormat.POLARS_SERIES,
        DataFormat.POLARS_LAZYFRAME,
        DataFormat.PYARROW_ARRAY,
    }:
        update_column_config(columns_config, INDEX_IDENTIFIER, {"hidden": True})


def _convert_column_config_to_json(column_config_mapping: ColumnConfigMapping) -> str:
    try:
        # Ignore all None values and prefix columns specified by numerical index:
        return json.dumps(
            {
                (f"{NUMERICAL_POSITION_PREFIX}{k!s}" if isinstance(k, int) else k): v
                for (k, v) in remove_none_values(column_config_mapping).items()
            },
            allow_nan=False,
        )
    except ValueError as ex:
        raise StreamlitAPIException(
            f"The provided column config cannot be serialized into JSON: {ex}"
        ) from ex


def marshall_column_config(
    proto: DataframeProto, column_config_mapping: ColumnConfigMapping
) -> None:
    """Marshall the column config into the Dataframe proto.

    Parameters
    ----------
    proto : DataframeProto
        The proto to marshall into.

    column_config_mapping : ColumnConfigMapping
        The column config to marshall.
    """

    proto.columns = _convert_column_config_to_json(column_config_mapping)
