"""
Plotly charts for TEI analyses.

Each function returns a ``plotly.graph_objects.Figure`` so callers can
``.show()`` (notebook), pass to ``st.plotly_chart`` (Streamlit), or write to
HTML / image. No styling is hard-coded beyond a neutral default palette.
"""

from __future__ import annotations

from collections.abc import Iterable

import plotly.graph_objects as go

PALETTE = {
    "benefits": "#2E7D32",  # green
    "costs": "#C62828",  # red
    "net_positive": "#1565C0",  # blue
    "net_negative": "#C62828",
    "cumulative": "#616161",  # grey
}

#: Visual theme — override per study/client with :func:`apply_theme`.
#: Hex colours; fonts are CSS font-family strings.
THEME = {
    "heading_font": "Helvetica Neue, Arial, sans-serif",
    "body_font": "Helvetica, Arial, sans-serif",
    "font_color": "#1F2937",
    # Circle-chart slice colours a–j, used in order.
    "pie_colors": [
        "#1565C0",  # a
        "#2E7D32",  # b
        "#C62828",  # c
        "#F9A825",  # d
        "#6A1B9A",  # e
        "#00838F",  # f
        "#EF6C00",  # g
        "#5D4037",  # h
        "#37474F",  # i
        "#AD1457",  # j
    ],
    "bar_green": "#2E7D32",
    "bar_red": "#C62828",
}


def apply_theme(**overrides) -> dict:
    """
    Override theme values for all charts in this session.

    Accepts any THEME key. ``pie_colors`` may be a list (used in order) or
    a dict keyed ``"a"``–``"j"`` (sorted alphabetically). Returns the
    active theme. Example::

        from core.notebook_helpers import charts
        charts.apply_theme(
            heading_font="Georgia, serif",
            font_color="#102A43",
            pie_colors={"a": "#1565C0", "b": "#2E7D32"},
            bar_green="#1B5E20",
            bar_red="#B71C1C",
        )
    """
    for key, value in overrides.items():
        if key not in THEME:
            raise KeyError(
                f"Unknown theme key {key!r}. Valid keys: {sorted(THEME)}"
            )
        if key == "pie_colors" and isinstance(value, dict):
            value = [value[k] for k in sorted(value)]
        THEME[key] = value
    return THEME


def _themed(fig: go.Figure) -> go.Figure:
    """Apply theme fonts/colours to a figure's layout."""
    fig.update_layout(
        font={"family": THEME["body_font"], "color": THEME["font_color"]},
        title_font={
            "family": THEME["heading_font"],
            "color": THEME["font_color"],
        },
        legend_font={"family": THEME["body_font"], "color": THEME["font_color"]},
    )
    return fig


def cashflow_chart(
    yearly_breakdown: list[dict],
    *,
    title: str = "Cash Flow Analysis (Risk-Adjusted)",
    initial_cost: float = 0.0,
) -> go.Figure:
    """
    Stacked bars of benefits & costs by year + cumulative net line.

    Mirrors the chart on page 25 of the Forrester Amazon Connect TEI study.
    """
    if not yearly_breakdown:
        return go.Figure(layout={"title": title})

    years = ["Initial"] + [f"Year {row['year']}" for row in yearly_breakdown]
    benefits = [0.0] + [float(row.get("benefits", 0)) for row in yearly_breakdown]
    costs = [-float(initial_cost)] + [
        -float(row.get("costs", 0)) for row in yearly_breakdown
    ]
    # cumulative_net assumes initial cost has already been deducted
    cumulative = [-float(initial_cost)] + [
        float(row.get("cumulative_net", 0)) for row in yearly_breakdown
    ]

    fig = go.Figure()
    fig.add_bar(
        name="Total benefits",
        x=years,
        y=benefits,
        marker_color=THEME["bar_green"],
    )
    fig.add_bar(
        name="Total costs",
        x=years,
        y=costs,
        marker_color=THEME["bar_red"],
    )
    fig.add_scatter(
        name="Cumulative net benefits",
        x=years,
        y=cumulative,
        mode="lines+markers",
        line={"color": PALETTE["cumulative"], "width": 3},
    )
    fig.update_layout(
        title=title,
        barmode="relative",
        yaxis_tickformat="$,.0f",
        legend={"orientation": "h", "y": -0.15},
        margin={"l": 40, "r": 20, "t": 60, "b": 40},
    )
    return _themed(fig)


def benefits_bar(items: list[dict], *, title: str = "Benefits (Three-Year)") -> go.Figure:
    """Horizontal bars of risk-adjusted three-year totals per benefit."""
    labels: list[str] = []
    totals: list[float] = []
    for it in items:
        rf = float(it.get("risk_adjustment") or 0.0)
        yv = it.get("year_values") or {}
        ra_total = sum(float(v or 0) * (1.0 - rf) for v in yv.values())
        labels.append(it.get("label", "") or it.get("field_key", ""))
        totals.append(ra_total)

    fig = go.Figure(
        go.Bar(
            x=totals,
            y=labels,
            orientation="h",
            marker_color=THEME["bar_green"],
            text=[f"${t/1_000_000:,.1f}M" for t in totals],
            textposition="auto",
        )
    )
    fig.update_layout(
        title=title,
        xaxis_tickformat="$,.0f",
        yaxis={"autorange": "reversed"},
        margin={"l": 40, "r": 20, "t": 60, "b": 40},
    )
    return _themed(fig)


def cost_breakdown_pie(
    items: list[dict], *, title: str = "Cost Breakdown (Three-Year, Risk-Adjusted)"
) -> go.Figure:
    """Pie chart of risk-adjusted costs by category/label."""
    labels: list[str] = []
    values: list[float] = []
    for it in items:
        rf = float(it.get("risk_adjustment") or 0.0)
        yv = it.get("year_values") or {}
        initial = float(it.get("initial") or 0.0)
        ra_total = (
            initial * (1.0 + rf)
            + sum(float(v or 0) * (1.0 + rf) for v in yv.values())
        )
        labels.append(it.get("label", "") or it.get("field_key", ""))
        values.append(ra_total)

    fig = go.Figure(go.Pie(labels=labels, values=values, hole=0.35,
                       marker={"colors": THEME["pie_colors"]}))
    fig.update_layout(title=title, margin={"l": 40, "r": 20, "t": 60, "b": 40})
    return _themed(fig)


def benefits_vs_costs_by_year(
    benefit_items: list[dict],
    cost_items: list[dict],
    *,
    title: str = "Benefits vs Costs by Year (Risk-Adjusted)",
) -> go.Figure:
    """
    Grouped bars of risk-adjusted benefits and costs per year, with an
    Initial (Year 0) column for one-time costs.

    Accepts the friendly value rows from ``TEIClient.get_values``:
    benefit values are nominal (field-level risk adjustment applied here);
    cost values are stored already risk-adjusted (Palladium convention),
    with ``initial`` carrying the Year-0 amount.
    """
    years: set[int] = set()
    for it in [*benefit_items, *cost_items]:
        years.update(int(y) for y in (it.get("year_values") or {}))
    year_list = sorted(years) or [1, 2, 3]

    benefits_by_year: dict[int, float] = dict.fromkeys(year_list, 0.0)
    costs_by_year: dict[int, float] = dict.fromkeys(year_list, 0.0)
    initial_total = 0.0

    for it in benefit_items:
        rf = float(it.get("risk_adjustment") or 0.0)
        for y, v in (it.get("year_values") or {}).items():
            benefits_by_year[int(y)] += float(v or 0) * (1.0 - rf)
    for it in cost_items:
        initial_total += float(it.get("initial") or 0.0)
        for y, v in (it.get("year_values") or {}).items():
            costs_by_year[int(y)] += float(v or 0)

    x = ["Initial"] + [f"Year {y}" for y in year_list]
    benefits = [0.0] + [benefits_by_year[y] for y in year_list]
    costs = [initial_total] + [costs_by_year[y] for y in year_list]

    fig = go.Figure()
    fig.add_bar(name="Benefits", x=x, y=benefits, marker_color=THEME["bar_green"],
                text=[f"${v/1_000_000:,.1f}M" if v else "" for v in benefits],
                textposition="outside")
    fig.add_bar(name="Costs", x=x, y=costs, marker_color=THEME["bar_red"],
                text=[f"${v/1_000_000:,.1f}M" if v else "" for v in costs],
                textposition="outside")
    fig.update_layout(
        title=title,
        barmode="group",
        yaxis_tickformat="$,.0f",
        legend={"orientation": "h", "y": -0.15},
        margin={"l": 40, "r": 20, "t": 60, "b": 40},
    )
    return _themed(fig)


def scenario_comparison(scenarios: dict) -> go.Figure:
    """Grouped bars comparing NPV and Costs PV across scenarios."""
    keys: list[str] = list(scenarios.keys())
    if not keys:
        return go.Figure()
    benefits = [float(scenarios[k].get("total_benefits_pv") or 0) for k in keys]
    costs = [float(scenarios[k].get("total_costs_pv") or 0) for k in keys]
    npvs = [float(scenarios[k].get("npv") or 0) for k in keys]

    fig = go.Figure()
    fig.add_bar(name="Benefits PV", x=keys, y=benefits, marker_color=THEME["bar_green"])
    fig.add_bar(name="Costs PV", x=keys, y=costs, marker_color=THEME["bar_red"])
    fig.add_bar(name="NPV", x=keys, y=npvs, marker_color=PALETTE["net_positive"])
    fig.update_layout(
        title="Scenario Comparison",
        barmode="group",
        yaxis_tickformat="$,.0f",
        legend={"orientation": "h", "y": -0.15},
    )
    return _themed(fig)


def cumulative_benefits_chart(
    yearly_breakdown: list[dict],
    *,
    title: str = "Cumulative Net Benefits",
) -> go.Figure:
    """Single-line cumulative net benefits trajectory."""
    if not yearly_breakdown:
        return go.Figure(layout={"title": title})
    years = [f"Year {row['year']}" for row in yearly_breakdown]
    cumulative = [float(row.get("cumulative_net", 0)) for row in yearly_breakdown]
    fig = go.Figure(
        go.Scatter(
            x=years,
            y=cumulative,
            mode="lines+markers",
            fill="tozeroy",
            line={"color": PALETTE["net_positive"], "width": 3},
        )
    )
    fig.update_layout(title=title, yaxis_tickformat="$,.0f")
    return _themed(fig)


def waterfall(values: Iterable[tuple[str, float]], *, title: str = "TEI Waterfall") -> go.Figure:
    """
    Generic waterfall (pass tuples of (label, value)).

    Used by 03_business_case to show: Benefits PV → Costs PV → NPV.
    """
    labels, amounts = zip(*values, strict=True) if values else ([], [])
    measures = ["relative"] * (len(labels) - 1) + ["total"] if labels else []
    fig = go.Figure(
        go.Waterfall(
            x=list(labels),
            y=list(amounts),
            measure=measures,
            text=[f"${v/1_000_000:,.1f}M" for v in amounts],
            textposition="outside",
            increasing={"marker": {"color": THEME["bar_green"]}},
            decreasing={"marker": {"color": THEME["bar_red"]}},
            totals={"marker": {"color": PALETTE["net_positive"]}},
        )
    )
    fig.update_layout(title=title, yaxis_tickformat="$,.0f")
    return _themed(fig)