"""Reusable chart components for Trades analysis."""

from __future__ import annotations

from typing import TYPE_CHECKING

import polars as pl

from .types import (
    ENTRY_REASON_LABELS,
    EXIT_REASON_LABELS,
    AxisExpr,
    ColorBy,
)

if TYPE_CHECKING:
    import plotly.graph_objects as go

    from ..mae_mfe.config import ChartConfig


def _import_plotly() -> tuple:
    """Lazy import Plotly with helpful error message."""
    try:
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots

        return go, make_subplots
    except ImportError as e:
        raise ImportError(
            "Plotly is required for Trades analysis visualization. "
            "Install with: pip install 'tradepose-client[analysis]' "
            "or: pip install plotly>=6.3.1"
        ) from e


def resolve_axis(df: pl.DataFrame, axis: AxisExpr) -> tuple[str, pl.Series]:
    """Resolve column name or pl.Expr to axis label and values.

    Args:
        df: DataFrame containing the data.
        axis: Column name (str) or Polars expression.

    Returns:
        Tuple of (label, values) where label is the axis title and
        values is the computed Series. Null values are filled with 0.

    Example:
        >>> label, values = resolve_axis(df, "mae")
        >>> label, values = resolve_axis(df, pl.col("mae") / pl.col("mae_volatility"))
    """
    # Normalize to pl.Expr
    if isinstance(axis, str):
        label = axis
        expr = pl.col(axis)
    else:
        # pl.Expr - extract alias or generate readable label
        label = axis.meta.output_name()
        if label is None or label.startswith("literal"):
            label = str(axis).replace("col(", "").replace(")", "").replace('"', "")
        expr = axis

    # Compute with fill_null(0) to handle missing values
    result = df.select(expr.fill_null(0).alias("_computed"))["_computed"]
    return label, result


def get_color_values(
    df: pl.DataFrame,
    color_by: ColorBy,
    config: ChartConfig,
) -> tuple[pl.Series, dict[str, str], str]:
    """Get color values, color map, and legend title based on color_by option.

    Args:
        df: DataFrame containing the data.
        color_by: Coloring option.
        config: Chart configuration for colors.

    Returns:
        Tuple of (color_series, color_map, legend_title).
    """
    if color_by == "outcome":
        # Win/Lose based on pnl
        color_series = df.select(
            pl.when(pl.col("pnl") > 0).then(pl.lit("Win")).otherwise(pl.lit("Lose")).alias("_color")
        )["_color"]
        color_map = {
            "Win": config.color_winner,
            "Lose": config.color_loser,
        }
        return color_series, color_map, "Outcome"

    elif color_by == "direction":
        # Long/Short
        color_series = df.select(
            pl.when(pl.col("direction") == 1)
            .then(pl.lit("Long"))
            .otherwise(pl.lit("Short"))
            .alias("_color")
        )["_color"]
        color_map = {
            "Long": config.color_long,
            "Short": config.color_short,
        }
        return color_series, color_map, "Direction"

    elif color_by == "exit_reason":
        if "exit_reason" not in df.columns:
            raise ValueError("exit_reason column not found in DataFrame")
        # Map exit reason codes to labels
        color_series = df.select(
            pl.col("exit_reason")
            .replace_strict(EXIT_REASON_LABELS, default="Unknown")
            .alias("_color")
        )["_color"]
        # Generate color palette
        unique_reasons = color_series.unique().to_list()
        palette = _get_categorical_palette(len(unique_reasons))
        color_map = dict(zip(unique_reasons, palette))
        return color_series, color_map, "Exit Reason"

    elif color_by == "entry_reason":
        if "entry_reason" not in df.columns:
            raise ValueError("entry_reason column not found in DataFrame")
        color_series = df.select(
            pl.col("entry_reason")
            .replace_strict(ENTRY_REASON_LABELS, default="Unknown")
            .alias("_color")
        )["_color"]
        unique_reasons = color_series.unique().to_list()
        palette = _get_categorical_palette(len(unique_reasons))
        color_map = dict(zip(unique_reasons, palette))
        return color_series, color_map, "Entry Reason"

    elif color_by == "strategy":
        if "strategy_name" not in df.columns:
            raise ValueError("strategy_name column not found in DataFrame")
        color_series = df["strategy_name"]
        unique_strategies = color_series.unique().to_list()
        palette = _get_categorical_palette(len(unique_strategies))
        color_map = dict(zip(unique_strategies, palette))
        return color_series, color_map, "Strategy"

    elif color_by == "blueprint":
        if "blueprint_name" not in df.columns:
            raise ValueError("blueprint_name column not found in DataFrame")
        color_series = df["blueprint_name"]
        unique_blueprints = color_series.unique().to_list()
        palette = _get_categorical_palette(len(unique_blueprints))
        color_map = dict(zip(unique_blueprints, palette))
        return color_series, color_map, "Blueprint"

    else:
        raise ValueError(f"Unknown color_by option: {color_by}")


def _get_categorical_palette(n: int) -> list[str]:
    """Get a categorical color palette with n colors."""
    # Plotly qualitative palette
    base_colors = [
        "#636EFA",
        "#EF553B",
        "#00CC96",
        "#AB63FA",
        "#FFA15A",
        "#19D3F3",
        "#FF6692",
        "#B6E880",
        "#FF97FF",
        "#FECB52",
    ]
    if n <= len(base_colors):
        return base_colors[:n]
    # Repeat colors if needed
    return (base_colors * ((n // len(base_colors)) + 1))[:n]


def scatter_component(
    df: pl.DataFrame,
    *,
    x: AxisExpr,
    y: AxisExpr,
    color_by: ColorBy = "outcome",
    hover_cols: list[str] | None = None,
    config: ChartConfig,
    title: str | None = None,
) -> go.Figure:
    """Create a reusable scatter plot component with pl.Expr support.

    Args:
        df: DataFrame containing the data.
        x: X-axis column name or Polars expression.
        y: Y-axis column name or Polars expression.
        color_by: Coloring option (outcome, direction, strategy, etc.).
        hover_cols: Additional columns to show on hover.
        config: Chart configuration.
        title: Optional chart title.

    Returns:
        Plotly Figure.

    Example:
        >>> scatter_component(df, x="mae", y="mfe", config=config)
        >>> scatter_component(
        ...     df,
        ...     x=pl.col("mae") / pl.col("mae_volatility"),
        ...     y="pnl",
        ...     config=config,
        ... )
    """
    go, _ = _import_plotly()

    x_label, x_values = resolve_axis(df, x)
    y_label, y_values = resolve_axis(df, y)
    color_series, color_map, legend_title = get_color_values(df, color_by, config)

    # Build hover text
    hover_data = []
    default_hover = ["pnl", "direction"]
    cols_to_show = list(set(default_hover + (hover_cols or [])))
    for i in range(len(df)):
        parts = [f"{x_label}: {x_values[i]:.4f}", f"{y_label}: {y_values[i]:.4f}"]
        for col in cols_to_show:
            if col in df.columns:
                val = df[col][i]
                if isinstance(val, float):
                    parts.append(f"{col}: {val:.4f}")
                else:
                    parts.append(f"{col}: {val}")
        hover_data.append("<br>".join(parts))

    fig = go.Figure()

    # Plot each color group
    for category, color in color_map.items():
        mask = color_series == category
        indices = [i for i, m in enumerate(mask.to_list()) if m]
        if not indices:
            continue

        fig.add_trace(
            go.Scatter(
                x=[x_values[i] for i in indices],
                y=[y_values[i] for i in indices],
                mode="markers",
                name=str(category),
                marker=dict(
                    color=color,
                    size=config.marker_size,
                    opacity=config.marker_opacity,
                ),
                hovertext=[hover_data[i] for i in indices],
                hoverinfo="text",
            )
        )

    fig.update_layout(
        title=title or f"{y_label} vs {x_label}",
        xaxis_title=x_label,
        yaxis_title=y_label,
        template=config.theme,
        width=config.width,
        height=config.height,
        showlegend=config.show_legend,
        legend_title_text=legend_title,
    )

    return fig


def box_component(
    df: pl.DataFrame,
    *,
    metric: AxisExpr,
    group_by: str,
    show_outcome: bool = True,
    config: ChartConfig,
    title: str | None = None,
) -> go.Figure:
    """Create a reusable box plot component grouped by category.

    Args:
        df: DataFrame containing the data.
        metric: Metric column or expression to analyze.
        group_by: Column to group by (exit_reason, entry_reason, blueprint_name, etc.).
        show_outcome: If True, split boxes by Win/Lose within each group.
        config: Chart configuration.
        title: Optional chart title.

    Returns:
        Plotly Figure.
    """
    go, _ = _import_plotly()

    metric_label, metric_values = resolve_axis(df, metric)

    # Get group labels
    if group_by == "exit_reason":
        group_series = df.select(
            pl.col("exit_reason")
            .replace_strict(EXIT_REASON_LABELS, default="Unknown")
            .alias("_group")
        )["_group"]
    elif group_by == "entry_reason":
        group_series = df.select(
            pl.col("entry_reason")
            .replace_strict(ENTRY_REASON_LABELS, default="Unknown")
            .alias("_group")
        )["_group"]
    else:
        group_series = df[group_by]

    fig = go.Figure()

    unique_groups = sorted(set(group_series.to_list()))

    if show_outcome and "pnl" in df.columns:
        outcome_series = df.select(
            pl.when(pl.col("pnl") > 0)
            .then(pl.lit("Win"))
            .otherwise(pl.lit("Lose"))
            .alias("_outcome")
        )["_outcome"]

        for group in unique_groups:
            for outcome in ["Win", "Lose"]:
                mask = (group_series == group) & (outcome_series == outcome)
                indices = [i for i, m in enumerate(mask.to_list()) if m]
                if not indices:
                    continue

                color = config.color_winner if outcome == "Win" else config.color_loser
                fig.add_trace(
                    go.Box(
                        y=[metric_values[i] for i in indices],
                        name=f"{group} ({outcome})",
                        marker_color=color,
                        boxpoints="outliers",
                    )
                )
    else:
        palette = _get_categorical_palette(len(unique_groups))
        for i, group in enumerate(unique_groups):
            mask = group_series == group
            indices = [j for j, m in enumerate(mask.to_list()) if m]
            if not indices:
                continue

            fig.add_trace(
                go.Box(
                    y=[metric_values[j] for j in indices],
                    name=str(group),
                    marker_color=palette[i % len(palette)],
                    boxpoints="outliers",
                )
            )

    fig.update_layout(
        title=title or f"{metric_label} by {group_by}",
        yaxis_title=metric_label,
        template=config.theme,
        width=config.width,
        height=config.height,
        showlegend=config.show_legend,
    )

    return fig


def distribution_component(
    df: pl.DataFrame,
    *,
    metric: AxisExpr,
    by_outcome: bool = True,
    bins: int = 50,
    config: ChartConfig,
    title: str | None = None,
) -> go.Figure:
    """Create a distribution histogram component.

    Args:
        df: DataFrame containing the data.
        metric: Metric column or expression to analyze.
        by_outcome: If True, show separate histograms for Win/Lose.
        bins: Number of histogram bins.
        config: Chart configuration.
        title: Optional chart title.

    Returns:
        Plotly Figure.
    """
    go, _ = _import_plotly()

    metric_label, metric_values = resolve_axis(df, metric)

    fig = go.Figure()

    if by_outcome and "pnl" in df.columns:
        outcome_series = df.select(
            pl.when(pl.col("pnl") > 0)
            .then(pl.lit("Win"))
            .otherwise(pl.lit("Lose"))
            .alias("_outcome")
        )["_outcome"]

        for outcome in ["Win", "Lose"]:
            mask = outcome_series == outcome
            indices = [i for i, m in enumerate(mask.to_list()) if m]
            if not indices:
                continue

            color = config.color_winner if outcome == "Win" else config.color_loser
            fig.add_trace(
                go.Histogram(
                    x=[metric_values[i] for i in indices],
                    name=outcome,
                    marker_color=color,
                    opacity=0.7,
                    nbinsx=bins,
                )
            )
        fig.update_layout(barmode="overlay")
    else:
        fig.add_trace(
            go.Histogram(
                x=metric_values.to_list(),
                name=metric_label,
                marker_color=config.color_neutral,
                nbinsx=bins,
            )
        )

    fig.update_layout(
        title=title or f"Distribution of {metric_label}",
        xaxis_title=metric_label,
        yaxis_title="Count",
        template=config.theme,
        width=config.width,
        height=config.height,
        showlegend=by_outcome,
    )

    return fig


def timeline_component(
    df: pl.DataFrame,
    *,
    metric: AxisExpr,
    period: str = "month",
    aggregation: str = "sum",
    config: ChartConfig,
    title: str | None = None,
) -> go.Figure:
    """Create a timeline bar chart grouped by time period.

    Args:
        df: DataFrame containing the data.
        metric: Metric column or expression to aggregate.
        period: Time period for grouping (day, week, month, quarter, year).
        aggregation: Aggregation function (sum, mean, count).
        config: Chart configuration.
        title: Optional chart title.

    Returns:
        Plotly Figure.
    """
    go, _ = _import_plotly()

    if "exit_time" not in df.columns:
        raise ValueError("exit_time column required for timeline analysis")

    metric_label, _ = resolve_axis(df, metric)

    # Handle both str and pl.Expr for metric
    if isinstance(metric, str):
        metric_expr = pl.col(metric)
    else:
        metric_expr = metric

    # Create period grouping expression
    period_expr: pl.Expr
    if period == "day":
        period_expr = pl.col("exit_time").dt.strftime("%Y-%m-%d")
    elif period == "week":
        period_expr = pl.col("exit_time").dt.strftime("%Y-W%W")
    elif period == "month":
        period_expr = pl.col("exit_time").dt.strftime("%Y-%m")
    elif period == "quarter":
        period_expr = (
            pl.col("exit_time").dt.year().cast(pl.Utf8)
            + "-Q"
            + pl.col("exit_time").dt.quarter().cast(pl.Utf8)
        )
    elif period == "year":
        period_expr = pl.col("exit_time").dt.year().cast(pl.Utf8)
    else:
        raise ValueError(f"Unknown period: {period}")

    # Aggregate by period
    agg_expr: pl.Expr
    if aggregation == "sum":
        agg_expr = metric_expr.sum()
    elif aggregation == "mean":
        agg_expr = metric_expr.mean()
    elif aggregation == "count":
        agg_expr = pl.count()
    else:
        raise ValueError(f"Unknown aggregation: {aggregation}")

    grouped = (
        df.with_columns(period_expr.alias("_period"))
        .group_by("_period")
        .agg(agg_expr.alias("_value"))
        .sort("_period")
    )

    periods = grouped["_period"].to_list()
    values = grouped["_value"].to_list()

    # Color bars by positive/negative
    colors = [config.color_winner if v >= 0 else config.color_loser for v in values]

    fig = go.Figure()
    fig.add_trace(
        go.Bar(
            x=periods,
            y=values,
            marker_color=colors,
            name=metric_label,
        )
    )

    agg_label = aggregation.title()
    fig.update_layout(
        title=title or f"{agg_label} {metric_label} by {period.title()}",
        xaxis_title=period.title(),
        yaxis_title=f"{agg_label} {metric_label}",
        template=config.theme,
        width=config.width,
        height=config.height,
    )

    return fig


def equity_curve_component(
    df: pl.DataFrame,
    *,
    metric_col: str = "pnl",
    ohlcv: pl.DataFrame | None = None,
    config: "ChartConfig",
    title: str | None = None,
) -> "go.Figure":
    """Create equity curve chart with optional benchmark comparison.

    Args:
        df: DataFrame containing trade data with exit_time and metric column.
        metric_col: Column to use for cumsum ("pnl" or "pnl_pct").
        ohlcv: Optional OHLCV DataFrame with 'ts' and 'open' columns for benchmark.
        config: Chart configuration.
        title: Optional chart title.

    Returns:
        Plotly Figure with equity curve(s).

    Example:
        >>> equity_curve_component(trades_df, config=config)
        >>> equity_curve_component(trades_df, metric_col="pnl_pct", ohlcv=ohlcv_df, config=config)
    """
    go, make_subplots = _import_plotly()

    if "exit_time" not in df.columns:
        raise ValueError("exit_time column required for equity curve")
    if metric_col not in df.columns:
        raise ValueError(f"{metric_col} column not found in DataFrame")

    # Sort by exit_time and calculate cumsum
    trades_sorted = df.sort("exit_time").select(
        [
            pl.col("exit_time"),
            pl.col(metric_col).cum_sum().alias("cum_pnl"),
        ]
    )

    is_pct = metric_col == "pnl_pct"
    y_label = "Cumulative Return (%)" if is_pct else "Cumulative PnL"

    # Determine if we need dual y-axis (when comparing with benchmark price)
    has_benchmark = ohlcv is not None
    show_benchmark_price = has_benchmark and not is_pct

    if show_benchmark_price:
        # Dual y-axis: left for PnL, right for price
        fig = make_subplots(specs=[[{"secondary_y": True}]])
    else:
        fig = go.Figure()

    # Add strategy equity curve
    fig.add_trace(
        go.Scatter(
            x=trades_sorted["exit_time"].to_list(),
            y=trades_sorted["cum_pnl"].to_list(),
            mode="lines",
            name="Strategy",
            line=dict(color=config.color_winner, width=2),
        ),
        secondary_y=False if show_benchmark_price else None,
    )

    # Add benchmark if OHLCV provided
    if has_benchmark:
        if "ts" not in ohlcv.columns:
            raise ValueError("OHLCV DataFrame must have 'ts' column")
        if "open" not in ohlcv.columns:
            raise ValueError("OHLCV DataFrame must have 'open' column")

        # Prepare benchmark data
        ohlcv_sorted = ohlcv.sort("ts")

        if is_pct:
            # Calculate benchmark cumulative return (pct_change cumsum)
            benchmark = ohlcv_sorted.select(
                [
                    pl.col("ts"),
                    (pl.col("open").pct_change().fill_null(0).cum_sum()).alias("cum_return"),
                ]
            )

            # Left join trades to benchmark and ffill
            combined = (
                benchmark.join(
                    trades_sorted.rename({"exit_time": "ts"}),
                    on="ts",
                    how="left",
                )
                .with_columns(pl.col("cum_pnl").forward_fill())
                .fill_null(0)
            )

            # Plot benchmark return
            fig.add_trace(
                go.Scatter(
                    x=combined["ts"].to_list(),
                    y=combined["cum_return"].to_list(),
                    mode="lines",
                    name="Benchmark",
                    line=dict(color=config.color_neutral, width=1.5, dash="dash"),
                )
            )

            # Plot strategy on OHLCV timeline (with ffill)
            fig.add_trace(
                go.Scatter(
                    x=combined["ts"].to_list(),
                    y=combined["cum_pnl"].to_list(),
                    mode="lines",
                    name="Strategy (aligned)",
                    line=dict(color=config.color_winner, width=2),
                    visible="legendonly",  # Hidden by default, show original
                )
            )
        else:
            # For absolute PnL, show benchmark price on secondary y-axis
            # Calculate benchmark PnL (price diff cumsum)
            benchmark = ohlcv_sorted.select(
                [
                    pl.col("ts"),
                    pl.col("open"),
                    (pl.col("open").diff().fill_null(0).cum_sum()).alias("cum_diff"),
                ]
            )

            # Left join trades to benchmark and ffill
            combined = (
                benchmark.join(
                    trades_sorted.rename({"exit_time": "ts"}),
                    on="ts",
                    how="left",
                )
                .with_columns(pl.col("cum_pnl").forward_fill())
                .fill_null(0)
            )

            # Plot benchmark price on secondary y-axis
            fig.add_trace(
                go.Scatter(
                    x=combined["ts"].to_list(),
                    y=combined["open"].to_list(),
                    mode="lines",
                    name="Benchmark Price",
                    line=dict(color=config.color_neutral, width=1, dash="dot"),
                ),
                secondary_y=True,
            )

            # Plot benchmark cumulative diff
            fig.add_trace(
                go.Scatter(
                    x=combined["ts"].to_list(),
                    y=combined["cum_diff"].to_list(),
                    mode="lines",
                    name="Benchmark PnL",
                    line=dict(color=config.color_loser, width=1.5, dash="dash"),
                ),
                secondary_y=False,
            )

            # Plot strategy on OHLCV timeline (with ffill)
            fig.add_trace(
                go.Scatter(
                    x=combined["ts"].to_list(),
                    y=combined["cum_pnl"].to_list(),
                    mode="lines",
                    name="Strategy (aligned)",
                    line=dict(color=config.color_winner, width=2),
                    visible="legendonly",
                ),
                secondary_y=False,
            )

    # Layout
    default_title = "Equity Curve" if not has_benchmark else "Strategy vs Benchmark"
    if is_pct:
        default_title += " (%)"

    if show_benchmark_price:
        fig.update_layout(
            title=title or default_title,
            xaxis_title="Time",
            template=config.theme,
            width=config.width,
            height=config.height,
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        )
        fig.update_yaxes(title_text=y_label, secondary_y=False)
        fig.update_yaxes(title_text="Price", secondary_y=True)
    else:
        fig.update_layout(
            title=title or default_title,
            xaxis_title="Time",
            yaxis_title=y_label,
            template=config.theme,
            width=config.width,
            height=config.height,
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        )

    return fig
