"""
himalayas/plot/renderers/cluster_labels
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
"""

from __future__ import annotations

from typing import (
    Any,
    Optional,
    Dict,
    Tuple,
    List,
    Sequence,
    TypedDict,
    Callable,
    TYPE_CHECKING,
)

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from ._cluster_label_types import ClusterLabelStats
from ._label_format import (
    apply_label_text_policy,
    collect_label_stats,
    compose_label_text,
    format_label_prefix,
)
from ._text_style import apply_text_style

if TYPE_CHECKING:
    from ..style import StyleConfig
    from ..track_layout import TrackLayoutManager
    from ...core.layout import ClusterLayout
    from ...core.matrix import Matrix


class TrackSpec(TypedDict, total=False):
    """
    Typed dictionary of resolved label track specifications.
    """

    name: str
    kind: str
    renderer: Callable[..., None]
    left_pad: float
    width: float
    right_pad: float
    enabled: bool
    payload: Dict[str, Any]
    x0: float
    x1: float


def _resolve_labels_and_layout(
    df: pd.DataFrame,
    kwargs: Dict[str, Any],
    fig: plt.Figure,
    matrix: Matrix,
    layout: ClusterLayout,
    style: StyleConfig,
    track_layout: TrackLayoutManager,
) -> Tuple[
    plt.Axes,
    float,
    List[TrackSpec],
    List[Tuple[int, int, int]],
    Dict[int, int],
    Dict[int, ClusterLabelStats],
    Dict[int, str],
    float,
    float,
    str,
    float,
    Optional[Tuple[str, ...]],
    bool,
    Optional[str],
]:
    """
    Resolves label data, overrides, axis layout, and text styling.

    Args:
        df (pd.DataFrame): Cluster label table with 'cluster', 'label', and optional
            'pval', 'qval', 'score', and 'fe'.
        kwargs (Dict[str, Any]): Renderer keyword arguments.
        fig (plt.Figure): Target figure.
        matrix (Matrix): Matrix object providing row count.
        layout (ClusterLayout): Cluster layout providing `cluster_spans`.
        style (StyleConfig): Style configuration.
        track_layout (TrackLayoutManager): Track layout manager.

    Returns:
        Tuple containing:
            - plt.Axes: Target label axis.
            - float: X-position for label text.
            - List[TrackSpec]: Resolved track specifications.
            - List[Tuple[int, int, int]]: Iterable of (cluster_id, start, end).
            - Dict[int, int]: Mapping cluster_id -> size.
            - Dict[int, ClusterLabelStats]: Mapping cluster_id -> (label, pval, qval, score, fe).
            - Dict[int, str]: Mapping cluster_id -> validated override label.
            - float: Minimum x-position for separator lines.
            - float: Maximum x-position for separator lines.
            - str: Font name for label text.
            - float: Font size for label text.
            - Optional[Tuple[str, ...]]: Fields to display in labels.
            - bool: Whether to skip unlabeled clusters.
            - Optional[str]: Label prefix mode.

    Raises:
        TypeError: If inputs have invalid types.
        ValueError: If required columns are missing or invalid.
    """
    # Validate cluster_labels DataFrame.
    if not isinstance(df, pd.DataFrame):
        raise TypeError("cluster_labels must be a pandas DataFrame.")
    if "cluster" not in df.columns or "label" not in df.columns:
        raise ValueError("cluster_labels DataFrame must contain columns: 'cluster', 'label'.")

    # Parse overrides and build label map.
    overrides = kwargs.get("overrides", None)
    override_map = _parse_label_overrides(overrides)
    label_map = _build_label_map(df, override_map)

    # Resolve spans, sizes, and label axis layout.
    spans = layout.cluster_spans
    cluster_sizes = layout.cluster_sizes
    ax_lab, label_text_x, tracks = _setup_label_axis(fig, matrix, style, track_layout)
    # Resolve separator line positions.
    sep_xmin = kwargs.get("label_sep_xmin", style.get("label_sep_xmin"))
    sep_xmax = kwargs.get("label_sep_xmax", style.get("label_sep_xmax"))
    if sep_xmin is None:
        sep_xmin = label_text_x
    if sep_xmax is None:
        sep_xmax = 1.0
    sep_xmin = float(np.clip(sep_xmin, 0.0, 1.0))
    sep_xmax = float(np.clip(sep_xmax, 0.0, 1.0))
    if sep_xmin > sep_xmax:
        sep_xmin, sep_xmax = sep_xmax, sep_xmin
    # Resolve text style options.
    font = kwargs.get("font", "Helvetica")
    fontsize = kwargs.get("fontsize", style.get("label_fontsize", 9))
    skip_unlabeled = kwargs.get("skip_unlabeled", False)
    label_fields = kwargs.get("label_fields", style["label_fields"])
    label_prefix = kwargs.get("label_prefix", None)

    # Validate label_fields
    if label_fields is not None and not isinstance(label_fields, (list, tuple)):
        raise TypeError("label_fields must be None or a list/tuple of strings")
    allowed_fields = {"label", "n", "p", "q", "fe"}
    if label_fields is not None and any(f not in allowed_fields for f in label_fields):
        raise ValueError(f"label_fields may only contain {allowed_fields}")
    if label_prefix not in {None, "cid", "alpha"}:
        raise ValueError("label_prefix must be one of {None, 'cid', 'alpha'}")
    resolved_label_fields = None if label_fields is None else tuple(label_fields)

    return (
        ax_lab,
        label_text_x,
        tracks,
        spans,
        cluster_sizes,
        label_map,
        override_map,
        sep_xmin,
        sep_xmax,
        font,
        fontsize,
        resolved_label_fields,
        bool(skip_unlabeled),
        label_prefix,
    )


def _render_tracks(
    ax_lab: plt.Axes,
    tracks: List[TrackSpec],
    *,
    matrix: Matrix,
    row_order: np.ndarray,
    spans: Sequence[Tuple[int, int, int]],
    label_map: Dict[int, ClusterLabelStats],
    style: StyleConfig,
    bar_labels_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
    """
    Renders all row-level and cluster-level tracks, and optional bar titles.

    Args:
        ax_lab (plt.Axes): Target label axis.
        tracks (List[TrackSpec]): List of track specifications.

    Kwargs:
        matrix (Matrix): Data matrix.
        row_order (np.ndarray): Row ordering indices.
        spans (Sequence[Tuple[int, int, int]]): Iterable of (cluster_id, start, end).
        label_map (Dict[int, ClusterLabelStats]): Mapping cluster_id -> (label, pval, qval, score, fe).
        style (StyleConfig): Style configuration.
        bar_labels_kwargs (Optional[Dict[str, Any]]): Bar title rendering options. Defaults to None.
    """
    # Render track content: data tracks and cluster-level tracks.
    for track in tracks:
        if track["kind"] == "row":
            track["renderer"](
                ax_lab,
                track["x0"],
                track["width"],
                track["payload"],
                matrix,
                row_order,
                style,
            )
    for track in tracks:
        if track["kind"] == "cluster":
            track["renderer"](
                ax_lab,
                track["x0"],
                track["width"],
                track["payload"],
                spans,
                label_map,
                style,
            )

    # Render optional bar titles.
    if bar_labels_kwargs is None:
        return

    # Render bar titles beneath tracks.
    bar_pad_pts = bar_labels_kwargs.get("pad", 2)
    bar_rotation = bar_labels_kwargs.get("rotation", 0)
    for track in tracks:
        title = track.get("payload", {}).get("title", None)
        if not title:
            continue
        x_center = (track.get("x0", 0.0) + track.get("x1", 0.0)) / 2.0
        txt = ax_lab.annotate(
            title,
            xy=(x_center, 0.0),
            xycoords=ax_lab.transAxes,
            xytext=(0, -bar_pad_pts),
            textcoords="offset points",
            ha="center",
            va="top",
            rotation=bar_rotation,
            clip_on=False,
        )
        apply_text_style(
            txt,
            font=bar_labels_kwargs.get("font", "Helvetica"),
            fontsize=bar_labels_kwargs.get("fontsize", 10),
            color=bar_labels_kwargs.get("color", style.get("text_color", "black")),
            alpha=bar_labels_kwargs.get("alpha", 1.0),
        )


def _render_cluster_text_and_separators(
    ax_lab: plt.Axes,
    *,
    spans: Sequence[Tuple[int, int, int]],
    cluster_sizes: Dict[int, int],
    label_map: Dict[int, ClusterLabelStats],
    override_map: Dict[int, str],
    label_text_x: float,
    sep_xmin: float,
    sep_xmax: float,
    font: str,
    fontsize: float,
    label_fields: Optional[Tuple[str, ...]],
    label_prefix: Optional[str],
    skip_unlabeled: bool,
    kwargs: Dict[str, Any],
    style: StyleConfig,
) -> None:
    """
    Renders cluster labels and separator lines.

    Args:
        ax_lab (plt.Axes): Target label axis.

    Kwargs:
        spans (Sequence[Tuple[int, int, int]]): Iterable of (cluster_id, start, end).
        cluster_sizes (Dict[int, int]): Mapping cluster_id -> size.
        label_map (Dict[int, ClusterLabelStats]): Mapping cluster_id -> (label, pval, qval, score, fe).
        override_map (Dict[int, str]): Mapping cluster_id -> validated override label.
        label_text_x (float): X-position for label text.
        sep_xmin (float): Minimum x-position for separator lines.
        sep_xmax (float): Maximum x-position for separator lines.
        font (str): Font name for label text.
        fontsize (float): Font size for label text.
        label_fields (Optional[Tuple[str, ...]]): Fields to display in labels.
        label_prefix (Optional[str]): Label prefix mode.
        skip_unlabeled (bool): Whether to skip unlabeled clusters.
        kwargs (Dict[str, Any]): Additional rendering options.
        style (StyleConfig): Style configuration.
    """
    max_words = kwargs.get("max_words", None)
    omit_words = kwargs.get("omit_words", style.get("label_omit_words", None))
    wrap_text = kwargs.get("wrap_text", True)
    wrap_width = kwargs.get("wrap_width", style.get("label_wrap_width", None))
    overflow = kwargs.get("overflow", "wrap")
    for cid, s, e in spans:
        y_center = (s + e) / 2.0
        # Choose placeholder or formatted label text for the cluster.
        if cid not in label_map:
            if skip_unlabeled:
                continue
            if label_fields is None and label_prefix is None:
                continue
            text = kwargs.get("placeholder_text", style["placeholder_text"])
            text_color = kwargs.get(
                "placeholder_color", kwargs.get("color", style["placeholder_color"])
            )
            text_alpha = kwargs.get(
                "placeholder_alpha", kwargs.get("alpha", style["placeholder_alpha"])
            )
        else:
            label, pval, qval, _score, fe = label_map[cid]
            is_override = cid in override_map
            prefix_active = label_prefix in {"cid", "alpha"} and not is_override
            force_label = prefix_active or is_override
            if (label_fields is None or "label" not in label_fields) and not is_override:
                label = ""
            if prefix_active:
                prefix = format_label_prefix(label_prefix, cid)
                label = f"{prefix} {label}" if label else prefix
            n_members = cluster_sizes.get(cid, None)
            text = _format_cluster_label(
                label,
                pval,
                qval,
                fe,
                n_members,
                label_fields=label_fields,
                force_label=force_label,
                max_words=max_words,
                omit_words=omit_words,
                wrap_text=wrap_text,
                wrap_width=wrap_width,
                overflow=overflow,
            )
            text_color = kwargs.get("color", style.get("text_color", "black"))
            text_alpha = kwargs.get("alpha", 0.9)
        # Draw label text and optional separator line.
        txt = ax_lab.text(
            label_text_x,
            y_center,
            text,
            va="center",
            ha="left",
            fontweight="normal",
            clip_on=False,
        )
        apply_text_style(txt, font=font, fontsize=fontsize, color=text_color, alpha=text_alpha)
        if s > 0:
            sep_color = kwargs.get("label_sep_color", style["label_sep_color"])
            sep_lw = kwargs.get("label_sep_lw", style["label_sep_lw"])
            sep_alpha = kwargs.get("label_sep_alpha", style["label_sep_alpha"])
            ax_lab.axhline(
                s - 0.5,
                xmin=sep_xmin,
                xmax=sep_xmax,
                color=sep_color,
                linewidth=sep_lw,
                alpha=sep_alpha,
                zorder=0,
            )


def _parse_label_overrides(
    overrides: Optional[Dict[int, str]] = None,
) -> Dict[int, str]:
    """
    Normalizes and validates per-cluster label overrides.

    Args:
        overrides (Dict[int, str] | None): Mapping cluster_id -> label string.
            Defaults to None.

    Returns:
        Dict[int, str]: Normalized override map keyed by cluster id.

    Raises:
        TypeError: If overrides or entries have invalid types.
    """
    if overrides is None:
        return {}
    # Validation
    if not isinstance(overrides, dict):
        raise TypeError("overrides must be a dict mapping cluster_id -> label string")

    # Normalize cluster ids and validate label strings.
    # Empty strings are allowed so callers can intentionally suppress label text.
    # for selected clusters without affecting bar tracks.
    override_map: Dict[int, str] = {}
    for key, value in overrides.items():
        cid = int(key)
        if not isinstance(value, str):
            raise TypeError("override values must be strings")
        override_map[cid] = value

    return override_map


def _build_label_map(
    df: pd.DataFrame,
    override_map: Dict[int, str],
) -> Dict[int, ClusterLabelStats]:
    """
    Resolves final label and p-value per cluster. Combines base labels from the DataFrame
    with any validated overrides.

    Args:
        df (pd.DataFrame): Cluster label table with 'cluster', 'label', and optional
            'pval', 'qval', 'score', and 'fe'.
        override_map (Dict[int, str]): Normalized overrides keyed by cluster id.

    Returns:
        Dict[int, ClusterLabelStats]: Mapping cluster id to (label, pval, qval, score, fe).

    Raises:
        ValueError: If overrides reference unknown cluster ids.
    """
    # Build base label map; overrides are label-only and must not alter stats.
    label_map: Dict[int, ClusterLabelStats] = {}
    for _, row in df.iterrows():
        cid = int(row["cluster"])
        base_label = str(row["label"])
        base_pval = row.get("pval", None)
        base_qval = row.get("qval", None)
        base_fe = row.get("fe", None)
        if "score" in row:
            base_score = row.get("score", None)
        elif "pval" in row:
            base_score = base_pval
        elif "qval" in row:
            base_score = base_qval
        else:
            base_score = None
        if cid in override_map:
            label = override_map[cid]
        else:
            label = base_label
        label_map[cid] = (label, base_pval, base_qval, base_score, base_fe)
    # Reject overrides that do not match any cluster id.
    if override_map:
        unknown = set(override_map) - set(label_map)
        if unknown:
            raise ValueError(
                "overrides contain cluster ids not present in cluster_labels: " f"{sorted(unknown)}"
            )

    return label_map


def _setup_label_axis(
    fig: plt.Figure,
    matrix: Matrix,
    style: StyleConfig,
    track_layout: TrackLayoutManager,
) -> Tuple[plt.Axes, float, List[TrackSpec]]:
    """
    Creates and configures the label axis and computes track layout. Initializes the label panel,
    draws the gutter, and resolves track x-positions.

    Args:
        fig (plt.Figure): Target figure.
        matrix (Matrix): Matrix object providing row count.
        style (StyleConfig): Style configuration.
        track_layout (TrackLayoutManager): Track layout manager.

    Returns:
        Tuple[plt.Axes, float, List[TrackSpec]]: (label axis, text x-position, resolved tracks).
    """
    n_rows = matrix.df.shape[0]
    # Set up label axis.
    label_axes = style["label_axes"]
    ax_lab = fig.add_axes(label_axes, frameon=False)
    ax_lab.set_xlim(0, 1)
    ax_lab.set_ylim(-0.5, n_rows - 0.5)  # align with matrix row indices.
    ax_lab.invert_yaxis()
    ax_lab.set_xticks([])
    ax_lab.set_yticks([])
    # Set up label gutter.
    gutter_w = style["label_gutter_width"]
    gutter_color = style["label_gutter_color"]
    ax_lab.add_patch(
        plt.Rectangle(
            (0.0, -0.5),
            gutter_w,
            n_rows,
            facecolor=gutter_color,
            edgecolor="none",
            zorder=0,
        )
    )
    # Compute track layout.
    label_text_pad = style.get("label_bar_pad", 0.01)
    base_x = style["label_x"]
    track_layout.compute_layout(base_x, gutter_w)
    tracks = track_layout.get_tracks()
    end_x = track_layout.get_end_x()
    if end_x is None:
        end_x = base_x + gutter_w
    label_text_x = end_x + label_text_pad

    return ax_lab, label_text_x, tracks


def _format_cluster_label(
    label: str,
    pval: Optional[float] = None,
    qval: Optional[float] = None,
    fe: Optional[float] = None,
    n_members: Optional[int] = None,
    *,
    label_fields: Optional[Tuple[str, ...]],
    force_label: bool = False,
    max_words: Optional[int] = None,
    omit_words: Optional[Sequence[str]] = None,
    wrap_text: bool = True,
    wrap_width: Optional[int] = None,
    overflow: str = "wrap",
) -> str:
    """
    Formats the cluster label text for display.

    Args:
        label (str): Base or overridden label.
        pval (float | None): P-value to display, if any. Defaults to None.
        qval (float | None): Q-value to display, if any. Defaults to None.
        fe (float | None): Fold enrichment to display, if any. Defaults to None.
        n_members (int | None): Cluster size. Defaults to None.

    Kwargs:
        label_fields (Optional[Tuple[str, ...]]): Fields to display.
        force_label (bool): Forces label text to render even when "label" is absent
            from label_fields. Defaults to False.
        max_words (Optional[int]): Maximum words to keep. Defaults to None.
        omit_words (Optional[Sequence[str]]): Words to omit (case-insensitive). Defaults to None.
        wrap_text (bool): Whether to wrap label text. Defaults to True.
        wrap_width (Optional[int]): Characters per wrapped line. Defaults to None.
        overflow (str): Truncation mode, one of {"wrap", "ellipsis"}. Defaults to "wrap".

    Returns:
        str: Final formatted label text.
    """
    # Assemble stats for requested fields.
    pval_value = pval if pval is not None and not pd.isna(pval) else None
    qval_value = qval if qval is not None and not pd.isna(qval) else None
    fe_value = fe if fe is not None and not pd.isna(fe) else None
    has_label, stats = collect_label_stats(
        label_fields,
        n_members=n_members,
        pval=pval_value,
        qval=qval_value,
        fe=fe_value,
        force_label=force_label,
    )

    # Apply label-only text policy, then append stats in a stable format.
    label_text = apply_label_text_policy(
        label,
        omit_words=omit_words,
        max_words=max_words,
        overflow=overflow,
        wrap_text=wrap_text,
        wrap_width=wrap_width,
    )
    if not has_label and not stats:
        return label_text
    return compose_label_text(
        label_text,
        has_label=has_label,
        stats=stats,
        wrap_text=wrap_text,
        wrap_width=wrap_width,
    )


class ClusterLabelsRenderer:
    """
    Class for rendering the cluster label panel with tracks and annotations.
    """

    def __init__(self, df: pd.DataFrame, **kwargs: Any) -> None:
        """
        Initializes the ClusterLabelsRenderer instance.

        Args:
            df (pd.DataFrame): Cluster label table.

        Kwargs:
            **kwargs: Renderer keyword arguments. Defaults to {}.
        """
        self.df = df
        self.kwargs = dict(kwargs)

    def render(
        self,
        fig: plt.Figure,
        matrix: Matrix,
        layout: ClusterLayout,
        style: StyleConfig,
        track_layout: TrackLayoutManager,
        bar_labels_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Renders the cluster label panel with tracks and annotations. Coordinates override
        handling, layout, track rendering, and label drawing.

        Args:
            fig (plt.Figure): Target figure.
            matrix (Matrix): Matrix object.
            layout (ClusterLayout): Cluster layout.
            style (StyleConfig): Style configuration.
            track_layout (TrackLayoutManager): Track layout manager.
            bar_labels_kwargs (Optional[Dict[str, Any]]): Bar title rendering options. Defaults to None.
        """
        df = self.df
        kwargs = self.kwargs
        # Resolve labels, layout, and styling.
        (
            ax_lab,
            label_text_x,
            tracks,
            spans,
            cluster_sizes,
            label_map,
            override_map,
            sep_xmin,
            sep_xmax,
            font,
            fontsize,
            label_fields,
            skip_unlabeled,
            label_prefix,
        ) = _resolve_labels_and_layout(
            df,
            kwargs,
            fig,
            matrix,
            layout,
            style,
            track_layout,
        )
        # Render tracks and cluster labels/separators.
        row_order = layout.leaf_order
        _render_tracks(
            ax_lab,
            tracks,
            matrix=matrix,
            row_order=row_order,
            spans=spans,
            label_map=label_map,
            style=style,
            bar_labels_kwargs=bar_labels_kwargs,
        )
        _render_cluster_text_and_separators(
            ax_lab,
            spans=spans,
            cluster_sizes=cluster_sizes,
            label_map=label_map,
            override_map=override_map,
            label_text_x=label_text_x,
            sep_xmin=sep_xmin,
            sep_xmax=sep_xmax,
            font=font,
            fontsize=fontsize,
            label_fields=label_fields,
            label_prefix=label_prefix,
            skip_unlabeled=skip_unlabeled,
            kwargs=kwargs,
            style=style,
        )
