from typing import Literal
from itertools import cycle

import uuid

import plotly.graph_objects as go

from dash import html, dcc, callback, Input, Output, State

import polars as pl

from .loading_test import DatasetConfig
from .data_source import DataSource

GRAPH_COLORS = [
    "#1A9D49",
    "#075745",
    "#1D9DE2",
    "#0F2080",
    "#C78800",
    "#471F00",
    "#C775A7",
    "#A3136C",
    "#909090",
    "#000000",
]


class GraphDisplay(html.Div):
    """
    The class handles its own global state.
    This is the state that is passed to components downstream
    """

    class ids:
        series_store = lambda aio_id: {
            "component": "GraphDisplay",
            "subcomponent": "series-store",
            "aio_id": aio_id,
        }
        settings_store = lambda aio_id: {
            "component": "GraphDisplay",
            "subcomponent": "settings-store",
            "aio_id": aio_id,
        }
        graph = lambda aio_id: {
            "component": "GraphDisplay",
            "subcomponent": "graph",
            "aio_id": aio_id,
        }

    # Make the ids class a public class
    ids = ids

    def __init__(self, datasets: dict[str, DatasetConfig], aio_id: None | str = None):
        """
        Expects the datasets.
        Can provide an aio_id if necessary.
        """
        if aio_id is None:
            aio_id = str(uuid.uuid4())

        super().__init__(
            children=[
                dcc.Store(data=[], id=self.ids.series_store(aio_id)),
                dcc.Store(data=[], id=self.ids.settings_store(aio_id)),
                html.Div(
                    children=[
                        dcc.Graph(id=self.ids.graph(aio_id)),
                    ],
                    style={"width": "100%", "height": "100%"},
                ),
            ]
        )

        @callback(
            Output(self.ids.graph(aio_id), "figure"),
            Input(self.ids.series_store(aio_id), "data"),
            Input(self.ids.settings_store(aio_id), "data"),
            State(self.ids.graph(aio_id), "figure"),
        )
        def change_graph(
            series_data,
            settings: dict[str, str | Literal["none", "discrete"]],
            old_fig,
        ):
            # Callback that updates the graph based on changed series settings or
            # graph settings
            base_year: str | None = settings.get("base_year")
            convert_method: Literal["none", "discrete"] = settings.get(
                "convert"
            )  # pyright: ignore
            fig = go.Figure()
            fig.update_layout(
                template="simple_white",
                title="Graf",
                autosize=True,
            )

            # Cycle for colors to alternative between allowed colors
            GRAPH_CYCLE = cycle(GRAPH_COLORS)

            if series_data == None:
                return fig

            for item in series_data:
                col = item["col"]
                dataset = item["dataset"]
                filename = item["path"]
                # Aggregation settings are optional
                convert_function = item.get("aggregation", None)
                # Shortens filepaths for displaying
                short_filename = "/".join(filename.split("/")[-2:])
                dataset_config = datasets[dataset]
                dataset = DataSource(
                    filename, dataset_config.index_col, dataset_config.index_pattern
                )

                data = dataset.data

                if (base_year is not None) and (convert_method != "none"):
                    agg_type = dataset_config.agg_type
                    if agg_type is None:
                        agg_type = dataset_config.agg_type_by_col
                    if agg_type is None:
                        agg_type = {}

                    data = dataset.set_base_year(
                        data,
                        base_year.split("-")[0],
                        col,
                        dataset_config.groupby_col,
                        convert_method,
                        agg_type,
                    )

                if (dataset_config.groupby_col) and ("groupby" in item):
                    for selected_val in item["groupby"]:
                        subset = data.filter(
                            pl.col(dataset_config.groupby_col) == selected_val
                        )
                        if convert_function is not None:
                            if convert_function == "pct":
                                subset = subset.with_columns(pl.col(col).pct_change())

                        # Add a line to the plot
                        fig.add_trace(
                            go.Scatter(
                                x=subset[dataset_config.index_col],
                                y=subset[col],
                                name=f"{short_filename} - {col} - {selected_val}",
                                line=dict(color=next(GRAPH_CYCLE)),
                                mode="lines",
                            )
                        )

                else:
                    if (convert_function is not None) and ("id" in convert_function):
                        if convert_function["id"] == "pct":
                            data = data.with_columns(pl.col(col).pct_change())

                    # Add the line to the plot
                    fig.add_trace(
                        go.Scatter(
                            x=data[dataset_config.index_col],
                            y=data[col],
                            name=col,
                            line=dict(color=next(GRAPH_CYCLE)),
                            mode="lines",
                        )
                    )

            # Tries to keep the selected timerange between callbacks
            selected_time_config = {}
            try:
                selected_time_config["range"] = old_fig["layout"]["xaxis"]["range"]
                selected_time_config["autorange"] = old_fig["layout"]["xaxis"][
                    "autorange"
                ]
            except:
                pass
            
            # Creates shortcuts for selecting timeframes
            fig.update_layout(
                xaxis=dict(
                    rangeselector=dict(
                        buttons=list(
                            [
                                dict(
                                    count=1,
                                    label="1m",
                                    step="month",
                                    stepmode="backward",
                                ),
                                dict(
                                    count=6,
                                    label="6m",
                                    step="month",
                                    stepmode="backward",
                                ),
                                dict(
                                    count=1, label="YTD", step="year", stepmode="todate"
                                ),
                                dict(
                                    count=1,
                                    label="1y",
                                    step="year",
                                    stepmode="backward",
                                ),
                                dict(step="all"),
                            ]
                        )
                    ),
                    rangeslider=dict(visible=True),
                    type="date",
                    **selected_time_config,
                )
            )
            return fig
