from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field
from enum import Enum

import pandas as pd
import plotly.graph_objects as go
from plotly.graph_objects import Figure
from plotly.subplots import make_subplots

from .quotes import normalize_quotes


class Color(str, Enum):
    RED = "#ef5350"
    GREEN = "#26a69a"
    BLUE = "#42a5f5"
    NAVY = "#0A1172"
    ORANGE = "#ffa726"
    PURPLE = "#ab47bc"
    BLACK = "#000000"
    WHITE = "#ffffff"
    YELLOW = "#ffff00"
    CYAN = "#00ffff"
    MAGENTA = "#ff00ff"
    GRAY = "#9e9e9e"
    TEAL = "#069494"


DEFAULT_PALETTE = (
    Color.BLUE,
    Color.PURPLE,
    Color.ORANGE,
    Color.NAVY,
    Color.MAGENTA,
    Color.TEAL,
)


class CharType(str, Enum):
    LINE = "line"
    BAR = "bar"


@dataclass
class SeriesPlot:
    """
    A class representing a series to be plotted.

    Attributes
    ----------
    series : pd.Series
        The data series to be plotted. It should have a datetime index and numeric values.
    type : CharType, optional
        The type of chart to use for the series, by default CharType.LINE.
    panel : int, optional
        The panel number where the series should be plotted, by default 0.
    label : str | None, optional
        The label for the series, by default None.
    color : Color | str | list[Color | str] | None, optional
        The color of the series, by default None.
    params : dict, optional
        Additional parameters for the series, by default an empty dictionary.
    """

    series: pd.Series
    type: CharType = field(default=CharType.LINE)
    panel: int = field(default=0)
    secondary_y: bool = field(default=False)
    label: str | None = field(default=None)
    color: Color | str | list[Color | str] | None = field(default=None)
    params: dict = field(default_factory=dict)

    def __post_init__(self):
        if not isinstance(self.series, pd.Series):
            raise ValueError("The 'series' attribute must be a pandas Series.")
        if not isinstance(self.series.index, pd.DatetimeIndex):
            raise ValueError("The index of the 'series' must be a DatetimeIndex.")
        if not pd.api.types.is_numeric_dtype(self.series):
            raise ValueError("The values of the 'series' must be numeric.")


def _infer_offdays_from_index(idx: pd.DatetimeIndex) -> list[pd.Timestamp]:
    """Infer offdays from DatetimeIndex by finding gaps larger than 1 day."""
    dates = idx.normalize().unique()
    all_days = pd.date_range(start=dates.min(), end=dates.max(), freq="D")
    offdays = all_days.difference(dates)
    return list(offdays)


def _plot_ohlcv(fig: Figure, ohlcv: pd.DataFrame) -> None:
    """
    Internal function to plot OHLCV data on the given figure.

    Parameters
    ----------
    fig : Figure
        The Plotly Figure object where the OHLCV data will be plotted.
    ohlcv : pd.DataFrame
        The OHLCV DataFrame containing 'open', 'high', 'low', 'close' columns and a datetime index; 'volume' column is optional.
    """
    fig.add_trace(
        go.Candlestick(
            x=ohlcv.index,
            open=ohlcv["open"],
            high=ohlcv["high"],
            low=ohlcv["low"],
            close=ohlcv["close"],
            name="Price",
        ),
        row=1,
        col=1,
    )

    if "volume" in ohlcv.columns:
        up = ohlcv["close"] >= ohlcv["open"]
        colors = [Color.GREEN.value if u else Color.RED.value for u in up]

        fig.add_trace(
            go.Bar(
                x=ohlcv.index,
                y=ohlcv["volume"],
                marker_color=colors,
                name="Volume",
            ),
            row=2,
            col=1,
        )


def _plot_indicator(fig: Figure, ind: SeriesPlot) -> None:
    """
    Internal function to plot an indicator on the given figure.

    Parameters
    ----------
    fig : Figure
        The Plotly Figure object where the indicator will be plotted.
    ind : SeriesPlot
        The SeriesPlot object representing the indicator to be plotted.
    """
    trace_extras = {
        "secondary_y": ind.secondary_y or None,
    }
    # Remove None values from trace_extras
    trace_extras = {k: v for k, v in trace_extras.items() if v is not None}

    chart_fun = go.Scatter if ind.type == CharType.LINE else go.Bar
    chart_params = (
        {
            "x": ind.series.index,
            "y": ind.series.values,
            "name": ind.label,
            **ind.params,
        }
        | (
            {"line": dict(color=ind.color)}
            if ind.type == CharType.LINE
            else {"marker_color": ind.color}
            if ind.color
            else {}
        )
        | ({"mode": "lines"} if ind.type == CharType.LINE else {})
    )

    fig.add_trace(
        chart_fun(**chart_params),
        row=ind.panel + 1,
        col=1,
        **trace_extras,
    )


def plot(
    quotes: pd.DataFrame,
    title: str | None = None,
    indicators: Sequence[SeriesPlot] | None = None,
    palette: Sequence[Color | str] | None = None,
    show: bool = True,
    **kwargs,
) -> Figure:
    """
    Plot the given quotes and indicators.

    Parameters
    ----------
    quotes : pd.DataFrame
        The quotes to be plotted. Must contain 'open', 'high', 'low', 'close' columns and have a datetime index; 'volume' column is optional.
    title : Optional[str], optional
        The title of the plot, by default None.
    indicators : Optional[Sequence[SeriesPlot]], optional
        A sequence of SeriesPlot objects representing the indicators to be plotted, by default None.
    **kwargs
        Additional keyword arguments for customization.

    Returns
    -------
    Figure
        A Plotly Figure object representing the plot.
    """
    ohlcv = normalize_quotes(quotes)
    indicators = indicators or []
    palette = list(palette or DEFAULT_PALETTE)

    ohlcv_panels = 2 if "volume" in ohlcv.columns else 1
    total_panels = max((ind.panel for ind in indicators), default=0) + 1
    total_panels = max(total_panels, ohlcv_panels)

    specs = [[{}] for _ in range(total_panels)]
    for ind in indicators:
        if ind.secondary_y:
            specs[ind.panel][0]["secondary_y"] = True

    row_heights = [0.7, 0.3] + [0.2] * (total_panels - 2) if total_panels > 1 else [1.0]

    fig = make_subplots(
        rows=total_panels,
        cols=1,
        shared_xaxes=True,
        vertical_spacing=0.01,
        row_heights=row_heights,
        specs=specs,
    )

    _plot_ohlcv(fig, ohlcv)

    for ind in indicators:
        if ind.color is None:
            ind.color = palette[0]
            palette = palette[1:] + palette[:1]  # Rotate the palette
        _plot_indicator(fig, ind)

    fig.update_layout(
        title_text=title,
        title_x=0.5,
        xaxis_rangeslider_visible=False,
        hovermode="x unified",
        **kwargs,
    )

    # Add rangebreaks to skip non-trading days
    rangebreaks = []
    rangebreaks.append(dict(values=_infer_offdays_from_index(ohlcv.index)))  # type: ignore
    if rangebreaks:
        fig.update_xaxes(rangebreaks=rangebreaks)

    if show:
        fig.show()

    return fig
