"""Discrete, diverging, and sequential accessible colour palettes.

Mirrors the R ``a11y_palette*`` family. Discrete and diverging palettes
draw from hardcoded RColorBrewer tables in ``_palette_data.py``.
Sequential palettes resolve through matplotlib colormaps (cividis,
viridis, plasma).
"""

import warnings

from a11yviz._palette_data import brewer
from a11yviz._rules import load_rules


def palette(name: str = "dark2_8", n: int | None = None,
            bg: str | None = None) -> list[str]:
    """Return hex codes for a colour-vision-aware categorical palette.

    Built-in: ``dark2_8`` (default, RColorBrewer Dark2), ``set2_8``
    (RColorBrewer Set2), ``paired_12`` (RColorBrewer Paired), ``aaa_5``
    (custom AAA-on-white). Pass ``n`` to truncate from the start. Pass
    ``bg`` (``"white"`` or ``"dark"``) to warn when the palette's
    ``safe_on`` tag does not match.
    """
    spec = _spec(name, "palettes")
    cols = _resolve_discrete(spec)
    if bg is not None:
        _warn_if_unsafe(name, spec, bg)
    if n is None:
        return cols
    return cols[:min(n, len(cols))]


def palette_div(name: str = "rdbu") -> dict:
    """Return ``{"low", "mid", "high"}`` anchors for a diverging gradient.

    Built-in: ``rdbu`` (default), ``puor``, ``brbg``, ``rdbu_dual``,
    ``puor_dual``, ``brbg_dual``, ``coolwarm_aaa``. Pass the result to
    a colour interpolator to materialise an N-step gradient.
    """
    spec = _spec(name, "diverging")
    src = spec.get("source", "literal")
    if src == "literal":
        return {"low": spec["low"], "mid": spec["mid"], "high": spec["high"]}
    if src == "rcolorbrewer":
        cols = brewer[spec["palette"]]
        pos = spec.get("positions", [10, 6, 2])
        return {"low":  cols[pos[0] - 1],
                "mid":  cols[pos[1] - 1],
                "high": cols[pos[2] - 1]}
    raise ValueError(f"Unsupported source '{src}' for diverging palette")


def palette_seq(name: str = "cividis", n: int | None = None):
    """Return a sequential gradient spec, or ``n`` interpolated hex codes.

    Built-in: ``cividis`` (default), ``viridis``, ``plasma``. Without
    ``n``, returns ``{"option", "begin", "end", "direction"}``. With
    ``n``, returns a list of ``n`` hex codes (requires matplotlib).
    """
    spec = _spec(name, "sequential")
    out = {
        "option":    spec["option"],
        "begin":     spec.get("begin", 0),
        "end":       spec.get("end", 1),
        "direction": spec.get("direction", 1),
    }
    if n is None:
        return out
    return _interpolate_seq(out, n)


def palette_info(name: str = "dark2_8") -> dict:
    """Return source spec, resolved colours, and WCAG metadata for a palette.

    Looks up ``name`` across discrete, diverging, and sequential
    families. For discrete palettes, ``colors`` holds the resolved hex
    list. For continuous palettes, ``colors`` is the gradient spec.
    """
    found = _find(name)
    if found is None:
        raise ValueError(f"Unknown palette '{name}'")
    kind, spec = found
    resolvers = {
        "discrete":   lambda: _resolve_discrete(spec),
        "diverging":  lambda: palette_div(name),
        "sequential": lambda: palette_seq(name),
    }
    return {"name": name, "type": kind, "colors": resolvers[kind](), **spec}


def palette_list(type: str | None = None) -> list[dict]:
    """List discrete, diverging, and sequential palettes in one table.

    ``type`` filters to one family (``"discrete"``, ``"diverging"``,
    or ``"sequential"``). Each row carries ``name``, ``type``,
    ``source``, ``n``, ``safe_on``, ``purpose``. ``n`` is ``None`` for
    continuous palettes.
    """
    rules = load_rules()
    families = (("discrete", "palettes"),
                ("diverging", "diverging"),
                ("sequential", "sequential"))
    rows = [_row(nm, kind, p)
            for kind, key in families
            for nm, p in rules[key].items()]
    if type is None:
        return rows
    if type not in {"discrete", "diverging", "sequential"}:
        raise ValueError(
            f"type must be 'discrete', 'diverging', or 'sequential' (got {type!r})"
        )
    return [r for r in rows if r["type"] == type]


# helpers ---------------------------------------------------------------------


def _spec(name: str, family: str) -> dict:
    pals = load_rules()[family]
    if name not in pals:
        raise ValueError(
            f"Unknown {family} palette '{name}'. "
            f"Available: {', '.join(pals.keys())}"
        )
    return pals[name]


def _find(name: str):
    rules = load_rules()
    for kind, key in (("discrete", "palettes"),
                      ("diverging", "diverging"),
                      ("sequential", "sequential")):
        if name in rules[key]:
            return kind, rules[key][name]
    return None


def _resolve_discrete(spec: dict) -> list[str]:
    src = spec.get("source", "literal")
    if src == "literal":
        return list(spec["colors"])
    if src == "rcolorbrewer":
        return list(brewer[spec["palette"]][:spec["n_colors"]])
    raise ValueError(f"Unsupported source '{src}' for discrete palette")


def _warn_if_unsafe(name: str, spec: dict, bg: str) -> None:
    if bg not in {"white", "dark"}:
        raise ValueError(f"bg must be 'white' or 'dark' (got {bg!r})")
    safe = spec.get("safe_on", "white")
    if safe == "both" or safe == bg:
        return
    note = spec.get("notes", "")
    if safe == "labeled":
        msg = (f"Palette '{name}' is mixed-contrast -- only safe when each "
               f"fill carries a high-contrast text label. {note}")
    else:
        msg = (f"Palette '{name}' is tagged safe_on='{safe}' but bg='{bg}' "
               f"was requested. {note}")
    warnings.warn(msg, stacklevel=3)


def _row(name: str, kind: str, spec: dict) -> dict:
    n = len(_resolve_discrete(spec)) if kind == "discrete" else None
    return {
        "name":    name,
        "type":    kind,
        "source":  spec.get("source", "literal"),
        "n":       n,
        "safe_on": spec.get("safe_on"),
        "purpose": spec.get("purpose"),
    }


def _interpolate_seq(spec: dict, n: int) -> list[str]:
    try:
        import matplotlib.colors as mcolors
        import numpy as np
        from matplotlib import colormaps
    except ImportError as e:
        raise ImportError(
            "palette_seq(n=...) requires matplotlib. "
            "Install with: pip install a11yviz[matplotlib]"
        ) from e
    cmap = colormaps[spec["option"]]
    positions = np.linspace(spec["begin"], spec["end"], n)[::spec["direction"]]
    return [mcolors.to_hex(cmap(p)) for p in positions]
