from __future__ import annotations

import json
from collections.abc import Callable, Sequence, Set
from typing import (
    TYPE_CHECKING,
    Any,
    Literal,
)

from gradio_client.documentation import document

from gradio.components.base import Component
from gradio.components.button import Button
from gradio.data_classes import GradioModel
from gradio.events import Events
from gradio.i18n import I18nData
from gradio.utils import set_default_buttons

if TYPE_CHECKING:
    import pandas as pd

    from gradio.components import Timer


class PlotData(GradioModel):
    columns: list[str]
    data: list[list[Any]]
    datatypes: dict[str, Literal["quantitative", "nominal", "temporal"]]
    mark: str


class NativePlot(Component):
    """
    Creates a native Gradio plot component to display data from a pandas DataFrame. Supports interactivity and updates.

    Demos: native_plots
    """

    EVENTS = [Events.select, Events.double_click]

    def __init__(
        self,
        value: pd.DataFrame | Callable | None = None,
        x: str | None = None,
        y: str | None = None,
        *,
        color: str | None = None,
        title: str | None = None,
        x_title: str | None = None,
        y_title: str | None = None,
        color_title: str | None = None,
        x_bin: str | float | None = None,
        y_aggregate: Literal["sum", "mean", "median", "min", "max", "count"]
        | None = None,
        color_map: dict[str, str] | None = None,
        colors_in_legend: list[str] | None = None,
        x_lim: list[float | None] | None = None,
        y_lim: list[float | None] = None,
        x_label_angle: float = 0,
        y_label_angle: float = 0,
        x_axis_format: str | None = None,
        y_axis_format: str | None = None,
        x_axis_labels_visible: bool | Literal["hidden"] = True,
        caption: str | I18nData | None = None,
        sort: Literal["x", "y", "-x", "-y"] | list[str] | None = None,
        tooltip: Literal["axis", "none", "all"] | list[str] = "axis",
        height: int | None = None,
        label: str | I18nData | None = None,
        show_label: bool | None = None,
        container: bool = True,
        scale: int | None = None,
        min_width: int = 160,
        every: Timer | float | None = None,
        inputs: Component | Sequence[Component] | Set[Component] | None = None,
        visible: bool | Literal["hidden"] = True,
        elem_id: str | None = None,
        elem_classes: list[str] | str | None = None,
        render: bool = True,
        buttons: list[Literal["fullscreen", "export"] | Button] | None = None,
        key: int | str | tuple[int | str, ...] | None = None,
        preserved_by_key: list[str] | str | None = "value",
    ):
        """
        Parameters:
            value: The pandas dataframe containing the data to display in the plot.
            x: Column corresponding to the x axis. Column can be numeric, datetime, or string/category.
            y: Column corresponding to the y axis. Column must be numeric.
            color: Column corresponding to series, visualized by color. Column must be string/category.
            title: The title to display on top of the chart.
            x_title: The title given to the x axis. By default, uses the value of the x parameter.
            y_title: The title given to the y axis. By default, uses the value of the y parameter.
            color_title: The title given to the color legend. By default, uses the value of color parameter.
            x_bin: Grouping used to cluster x values. If x column is numeric, should be number to bin the x values. If x column is datetime, should be string such as "1h", "15m", "10s", using "s", "m", "h", "d" suffixes.
            y_aggregate: Aggregation function used to aggregate y values, used if x_bin is provided or x is a string/category. Must be one of "sum", "mean", "median", "min", "max".
            color_map: Mapping of series to color names or codes. For example, {"success": "green", "fail": "#FF8888"}.
            colors_in_legend: List containing column names of the series to show in the legend. By default, all series are shown.
            height: The height of the plot in pixels.
            x_lim: A tuple or list containing the limits for the x-axis, specified as [x_min, x_max]. To fix only one of these values, set the other to None, e.g. [0, None] to scale from 0 to the maximum value. If x column is datetime type, x_lim should be timestamps.
            y_lim: A tuple of list containing the limits for the y-axis, specified as [y_min, y_max]. To fix only one of these values, set the other to None, e.g. [0, None] to scale from 0 to the maximum to value.
            x_label_angle: The angle of the x-axis labels in degrees offset clockwise.
            y_label_angle: The angle of the y-axis labels in degrees offset clockwise.
            x_axis_format: A d3 format string for the x-axis labels (e.g., ".2e" for scientific notation, "~g" for general format). If None, uses Vega-Lite's default formatting.
            y_axis_format: A d3 format string for the y-axis labels (e.g., ".2e" for scientific notation, "~g" for general format). If None, uses Vega-Lite's default formatting.
            x_axis_labels_visible: Whether the x-axis labels should be visible. Can be hidden when many x-axis labels are present.
            caption: The (optional) caption to display below the plot.
            sort: The sorting order of the x values, if x column is type string/category. Can be "x", "y", "-x", "-y", or list of strings that represent the order of the categories.
            tooltip: The tooltip to display when hovering on a point. "axis" shows the values for the axis columns, "all" shows all column values, and "none" shows no tooltips. Can also provide a list of strings representing columns to show in the tooltip, which will be displayed along with axis values.
            height: The height of the plot in pixels.
            label: The (optional) label to display on the top left corner of the plot.
            show_label: Whether the label should be displayed.
            container: If True, will place the component in a container - providing some extra padding around the border.
            scale: relative size compared to adjacent Components. For example if Components A and B are in a Row, and A has scale=2, and B has scale=1, A will be twice as wide as B. Should be an integer. scale applies in Rows, and to top-level Components in Blocks where fill_height=True.
            min_width: minimum pixel width, will wrap if not sufficient screen space to satisfy this value. If a certain scale value results in this Component being narrower than min_width, the min_width parameter will be respected first.
            every: Continously calls `value` to recalculate it if `value` is a function (has no effect otherwise). Can provide a Timer whose tick resets `value`, or a float that provides the regular interval for the reset Timer.
            inputs: Components that are used as inputs to calculate `value` if `value` is a function (has no effect otherwise). `value` is recalculated any time the inputs change.
            visible: Whether the plot should be visible.
            elem_id: An optional string that is assigned as the id of this component in the HTML DOM. Can be used for targeting CSS styles.
            elem_classes: An optional list of strings that are assigned as the classes of this component in the HTML DOM. Can be used for targeting CSS styles.
            render: If False, component will not render be rendered in the Blocks context. Should be used if the intention is to assign event listeners now but render the component later.
            buttons: A list of buttons to show for the component. Valid options are "fullscreen", "export", or a gr.Button() instance. The "fullscreen" button allows the user to view the plot in fullscreen mode. The "export" button allows the user to export and download the current view of the plot as a PNG image. Custom gr.Button() instances will appear in the toolbar with their configured icon and/or label, and clicking them will trigger any .click() events registered on the button. By default, no buttons are shown.
            key: in a gr.render, Components with the same key across re-renders are treated as the same component, not a new component. Properties set in 'preserved_by_key' are not reset across a re-render.
            preserved_by_key: A list of parameters from this component's constructor. Inside a gr.render() function, if a component is re-rendered with the same key, these (and only these) parameters will be preserved in the UI (if they have been changed by the user or an event listener) instead of re-rendered based on the values provided during constructor.
        """
        self.x = x
        self.y = y
        self.color = color
        self.title = title
        self.x_title = x_title
        self.y_title = y_title
        self.color_title = color_title
        self.x_bin = x_bin
        self.y_aggregate = y_aggregate
        self.color_map = color_map
        self.x_lim = x_lim
        self.y_lim = y_lim
        self.x_label_angle = x_label_angle
        self.y_label_angle = y_label_angle
        self.x_axis_format = x_axis_format
        self.y_axis_format = y_axis_format
        self.x_axis_labels_visible = x_axis_labels_visible
        self.caption = caption
        self.sort = sort
        self.tooltip = tooltip
        self.height = height
        self.buttons = set_default_buttons(buttons, None)
        self.colors_in_legend = colors_in_legend

        if label is None and show_label is None:
            show_label = False
        super().__init__(
            value=value,
            label=label,
            show_label=show_label,
            container=container,
            scale=scale,
            min_width=min_width,
            visible=visible,
            elem_id=elem_id,
            elem_classes=elem_classes,
            render=render,
            key=key,
            preserved_by_key=preserved_by_key,
            every=every,
            inputs=inputs,
        )

    def get_block_name(self) -> str:
        return "nativeplot"

    def get_mark(self) -> str:
        return "native"

    def preprocess(self, payload: PlotData | None) -> PlotData | None:
        """
        Parameters:
            payload: The data to display in a line plot.
        Returns:
            The data to display in a line plot.
        """
        return payload

    def postprocess(self, value: pd.DataFrame | dict | None) -> PlotData | None:
        """
        Parameters:
            value: Expects a pandas DataFrame containing the data to display in the line plot. The DataFrame should contain at least two columns:
            - one for the x-axis (corresponding to this component's `x` argument)
            - one for the y-axis (corresponding to `y`).
        Returns:
            The data to display in a line plot, in the form of an AltairPlotData dataclass, which includes the plot information as a JSON string, as well as the type of plot (in this case, "line").
        """
        # if None or update
        if value is None or isinstance(value, dict):
            return value

        def get_simplified_type(dtype):
            import pandas as pd

            if pd.api.types.is_numeric_dtype(dtype):
                return "quantitative"
            elif pd.api.types.is_string_dtype(
                dtype
            ) or pd.api.types.is_categorical_dtype(dtype):
                return "nominal"
            elif pd.api.types.is_datetime64_any_dtype(dtype):
                return "temporal"
            else:
                raise ValueError(f"Unsupported data type: {dtype}")

        split_json = json.loads(value.to_json(orient="split", date_unit="ms"))
        datatypes = {
            col: get_simplified_type(value[col].dtype) for col in value.columns
        }
        return PlotData(
            columns=split_json["columns"],
            data=split_json["data"],
            datatypes=datatypes,
            mark=self.get_mark(),
        )

    def example_payload(self) -> Any:
        return None

    def example_value(self) -> Any:
        import pandas as pd

        return pd.DataFrame({self.x: [1, 2, 3], self.y: [4, 5, 6]})

    def api_info(self) -> dict[str, Any]:
        return {"type": {}, "description": "any valid json"}


@document()
class BarPlot(NativePlot):
    """
    Creates a bar plot component to display data from a pandas DataFrame.

    Demos: bar_plot_demo
    Guides: creating-plots, time-plots
    """

    def get_block_name(self) -> str:
        return "nativeplot"

    def get_mark(self) -> str:
        return "bar"


@document()
class LinePlot(NativePlot):
    """
    Creates a line plot component to display data from a pandas DataFrame.

    Demos: line_plot_demo
    Guides: creating-plots, connecting-to-a-database
    """

    def get_block_name(self) -> str:
        return "nativeplot"

    def get_mark(self) -> str:
        return "line"


@document()
class ScatterPlot(NativePlot):
    """
    Creates a scatter plot component to display data from a pandas DataFrame.

    Demos: scatter_plot_demo
    Guides: creating-plots
    """

    def get_block_name(self) -> str:
        return "nativeplot"

    def get_mark(self) -> str:
        return "point"
