#!/usr/bin/env python3
"""Figure-class router. suggest() returns a class, never a toolkit."""

from __future__ import annotations

NUMERIC = frozenset({"loss-curve", "bar", "scatter", "histogram", "numeric-plot"})
RASTER = frozenset({"table", "photo", "apparatus", "ui"})
TENSOR = frozenset({
    "axis", "shape", "transpose", "broadcast", "gather",
    "shard", "contraction", "face",
})
DERIVE = frozenset({"rewrite-figure", "cancel-visual", "subst-visual"})
FIG = frozenset({
    "architecture", "pipeline", "data-flow", "state", "time",
    "dependency", "what-eats-what", "argument-map",
})

CLASS_SIGNALS = {
    "numeric": NUMERIC,
    "raster": RASTER,
    "tensor": TENSOR,
    "derive": DERIVE,
    "fig": FIG,
}

VECTOR = frozenset({"superfig", "supertensor", "superderive", "matplotlib"})
SIBLING = frozenset({"superfig", "supertensor", "superderive"})


def classify(signals: set[str]) -> list[str]:
    s = set(signals)
    hit = [c for c, vocab in CLASS_SIGNALS.items() if s & vocab]
    return hit or ["none"]


def suggest(signals: set[str], *, phase: int = 1) -> str:
    """Return a class name, or 'SPLIT:a+b+...' in CLASS_SIGNALS order."""
    classes = classify(signals)
    if len(classes) > 1:
        return "SPLIT:" + "+".join(classes)
    return classes[0]


def accepted(cls: str, *, phase: int = 1) -> frozenset[str]:
    if cls == "numeric":
        return frozenset({"screenshot", "matplotlib"})
    if cls == "raster":
        return frozenset({"screenshot"})
    if cls == "tensor":
        return frozenset({"supertensor"})
    if cls == "fig":
        return frozenset({"superfig"})
    if cls == "derive":
        return frozenset({"align", "superderive"} if phase >= 2 else {"align"})
    if cls == "none":
        return frozenset({"none"})
    raise KeyError(cls)


def include_pdf(fig_id: str) -> str:
    return f"figures/{fig_id}/build/{fig_id}.pdf"


def include_png(fig_id: str) -> str:
    return f"figures/{fig_id}/orig.png"


def check_row(
    signals: set[str],
    toolkit: str,
    *,
    status: str = "planned",
    include: str | None = None,
    fig_id: str = "F1",
    phase: int = 1,
) -> str | None:
    """First matching code, or None. SP010 is not an include/status rule."""
    classes = classify(signals)
    if len(classes) > 1:
        return "SP011"
    if toolkit not in accepted(classes[0], phase=phase):
        return "SP010"
    if toolkit == "none" and status != "dropped":
        return "SP013"
    if toolkit in {"align", "none"} and include is not None:
        return "SP014"
    if status == "included" and toolkit in VECTOR and include != include_pdf(fig_id):
        return "SP015"
    if status == "included" and toolkit == "screenshot" and include != include_png(fig_id):
        return "SP016"
    return None


def phase_from_env() -> int:
    import os
    raw = os.environ.get("SUPERPAPER_PHASE", "1")
    try:
        return int(raw)
    except ValueError:
        return 1