"""Legend and panel-label helpers for consistent multi-panel figures."""

from __future__ import annotations

import warnings
from typing import Optional

from matplotlib.axes import Axes
from matplotlib.colors import to_rgba

from .colors import INK

_PANEL_LOCS = {
    "upper left": (0.02, 0.98, "left", "top"),
    "upper right": (0.98, 0.98, "right", "top"),
    "lower left": (0.02, 0.08, "left", "bottom"),
    "lower right": (0.98, 0.08, "right", "bottom"),
}


def style_legend(
    ax: Axes,
    loc: str = "outside right upper",
    frameon: bool = False,
    title: Optional[str] = None,
    **kwargs,
):
    """Add a figure-level legend, placed outside the axes by default.

    Builds on `ax`'s handles/labels (or explicit `handles=`/`labels=` kwargs)
    but attaches the legend to `ax`'s figure via `fig.legend(...)`, so it sits
    outside the plot area rather than overlapping the data.

    A `title` is strongly encouraged: an untitled legend floating outside the
    axes loses its visual link to what it's describing, e.g.
    `style_legend(ax, title="Series")`. Without one, this emits a warning
    rather than failing — the legend still renders.
    """
    if title is None:
        warnings.warn(
            "style_legend() called without a title — an outside legend reads better with one, "
            "e.g. style_legend(ax, title='Series').",
            stacklevel=2,
        )

    fig = ax.get_figure()
    assert fig is not None, "ax must be attached to a figure"
    handles = kwargs.pop("handles", None)
    labels = kwargs.pop("labels", None)
    if handles is None or labels is None:
        handles, labels = ax.get_legend_handles_labels()

    # matplotlib-stubs' `loc` Literal doesn't include the "outside ..." compound
    # locations matplotlib actually supports at runtime (e.g. "outside right upper").
    legend = fig.legend(handles, labels, loc=loc, frameon=frameon, title=title, **kwargs)  # ty: ignore[invalid-argument-type]
    if legend.get_title() is not None:
        legend.get_title().set_fontweight("bold")
    return legend


def panel_label(
    ax: Axes,
    label: str,
    loc: str = "lower right",
    fontweight: str = "bold",
    box: bool = True,
    **kwargs,
):
    """Add a panel label like "(a)" for multi-panel thesis/paper figures.

    Defaults to the bottom-right corner, colored to match the axes' xlabel,
    on a light-grey semi-transparent rounded box with a slim solid border.
    Pass `box=False` for bare text with no box.
    """
    try:
        x, y, ha, va = _PANEL_LOCS[loc]
    except KeyError as exc:
        raise ValueError(f"Unknown panel_label loc {loc!r}. Choose from {sorted(_PANEL_LOCS)}.") from exc

    kwargs.setdefault("color", ax.xaxis.label.get_color())
    if box:
        kwargs.setdefault(
            "bbox",
            dict(
                boxstyle="round,pad=0.3",
                facecolor=to_rgba(INK["gridline"], alpha=0.8),
                edgecolor=INK["baseline"],
                linewidth=0.8,
            ),
        )

    return ax.text(
        x,
        y,
        f"({label})",
        transform=ax.transAxes,
        ha=ha,
        va=va,
        fontweight=fontweight,
        **kwargs,
    )