#!/usr/bin/env python3
"""Static source checks for supertensor figures.

The TeX package owns geometry at render time; this linter catches source-level
escapes that TeX cannot see: mutable ledgers, hand placement, formula order,
raw rectangles, repeated anonymous dimensions, and per-row hue budget.
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path


DIM_RE = re.compile(r"\\stdim\{([^{}]+)\}\{([^{}]+)\}")
ROLE_RE = re.compile(r"\\stsetrole\{([^{}]+)\}\{([^{}]+)\}")
FACE_RE = re.compile(
    r"\\(stface|ststack|stindexface)(?:\[[^\]]*\])?"
    r"\{([^{}]+)\}\{([^{}]*)\}\{([^{}]+)\}\{([^{}]+)\}",
    re.DOTALL,
)
ROW_RE = re.compile(r"\\strow\{[^{}]+\}\{[^{}]+\}(.*?)\\strowend", re.DOTALL)
ROW_NAME_RE = re.compile(r"\\strow\{([^{}]+)\}\{[^{}]+\}")
ROLE_KEY_RE = re.compile(r"role\s*=\s*([A-Za-z0-9_-]+)")
CALLOUT_RE = re.compile(r"\\stcallout\{([^{}]+)\}\{[^{}]+\}\{([^{}]+)\}")


def strip_comments(source: str) -> str:
    return re.sub(r"(?<!\\)%.*$", "", source, flags=re.MULTILINE)


def line_of(source: str, offset: int) -> int:
    return source.count("\n", 0, offset) + 1


def directives(source: str) -> set[str]:
    found: set[str] = set()
    for match in re.finditer(r"^\s*%\s*supertensor-lint:\s*(.+)$", source, re.MULTILINE):
        found.update(item.strip() for item in match.group(1).split(","))
    return found


def check_unique(
    source: str, pattern: re.Pattern[str], kind: str, errors: list[str]
) -> dict[str, str]:
    values: dict[str, str] = {}
    for match in pattern.finditer(source):
        name, value = (part.strip() for part in match.groups())
        previous = values.get(name)
        if previous is not None and previous != value:
            errors.append(
                f"line {line_of(source, match.start())}: {kind} {name!r} changes "
                f"from {previous!r} to {value!r}"
            )
        else:
            values[name] = value
    return values


def lint(path: Path) -> list[str]:
    raw = path.read_text(encoding="utf-8")
    allowed = directives(raw)
    source = strip_comments(raw)
    errors: list[str] = []

    dimensions = check_unique(source, DIM_RE, "axis", errors)
    roles = check_unique(source, ROLE_RE, "role", errors)

    formula_positions = [m.start() for m in re.finditer(r"\\sttopformula\b", source)]
    row_ends = [m.start() for m in re.finditer(r"\\strowend\b", source)]
    if not formula_positions and "allow-missing-formula" not in allowed:
        errors.append("missing \\sttopformula (add an explicit lint exemption for galleries/tests)")
    elif formula_positions and row_ends and formula_positions[-1] < row_ends[-1]:
        errors.append(
            f"line {line_of(source, formula_positions[-1])}: \\sttopformula must follow the last \\strowend"
        )

    tracked = set(re.findall(r"\\sttrack\{([^{}]+)\}", source))
    raw_dimensions: dict[str, list[int]] = {}
    for match in FACE_RE.finditer(source):
        _macro, name, coordinate, rows, cols = match.groups()
        if coordinate.strip() and name not in tracked and "allow-absolute" not in allowed:
            errors.append(
                f"line {line_of(source, match.start())}: absolute object {name!r} has no \\sttrack{{{name}}}"
            )
        for value in (rows.strip(), cols.strip()):
            if value.isdigit() and value != "1":
                raw_dimensions.setdefault(value, []).append(line_of(source, match.start()))

    if "allow-repeated-raw-dim" not in allowed:
        for value, lines in sorted(raw_dimensions.items()):
            if len(lines) > 1:
                errors.append(
                    f"lines {', '.join(map(str, lines))}: raw dimension {value!r} is reused; declare a symbolic axis with \\stdim"
                )

    if "allow-raw-tikz" not in allowed:
        for match in re.finditer(
            r"\\(?:draw|fill|path)\b[^;]*\brectangle\b", source, re.DOTALL
        ):
            errors.append(
                f"line {line_of(source, match.start())}: hand-drawn rectangle bypasses \\stface"
            )

    # \stgroup is a block. An unclosed one silently swallows the rest of the
    # band into the group's fit list, which TeX only reports as an empty row.
    opened = len(re.findall(r"\\stgroup(?![A-Za-z])", source))
    closed = len(re.findall(r"\\stgroupend(?![A-Za-z])", source))
    if opened != closed:
        errors.append(
            f"{opened} \\stgroup vs {closed} \\stgroupend: every group must be closed"
        )

    # A callout hangs off a BAND. Anchored to a face it becomes the floating
    # commentary card between two operands that layout.md forbids. The default
    # budget is ONE callout for the whole figure: a card per band is already a
    # dashboard, and at thumbnail size the cards are unreadable anyway. The
    # allow-multiple-callouts directive relaxes this to one per band.
    band_names = set(ROW_NAME_RE.findall(source))
    callout_anchors: dict[str, str] = {}
    first_callout: str | None = None
    for match in CALLOUT_RE.finditer(source):
        name, anchor = (part.strip() for part in match.groups())
        line = line_of(source, match.start())
        if anchor not in band_names:
            errors.append(
                f"line {line}: callout {name!r} is anchored to {anchor!r}, which is not a "
                f"\\strow band; a card hanging off a single object reads as a step"
            )
        elif anchor in callout_anchors:
            errors.append(
                f"line {line}: callout {name!r} is the second card on band {anchor!r} "
                f"(after {callout_anchors[anchor]!r}); one aside per band"
            )
        elif first_callout is not None and "allow-multiple-callouts" not in allowed:
            errors.append(
                f"line {line}: callout {name!r} is the second card in the figure "
                f"(after {first_callout!r}); the budget is one callout per figure — "
                f"move the text to \\stmeaningbox, or add the "
                f"allow-multiple-callouts exemption"
            )
        callout_anchors.setdefault(anchor, name)
        if first_callout is None:
            first_callout = name

    for row_number, row in enumerate(ROW_RE.finditer(source), start=1):
        row_roles = set(ROLE_KEY_RE.findall(row.group(1)))
        colors = {roles.get(role, "<undeclared>") for role in row_roles}
        active = colors - {"stGray"}
        if len(active) > 4:
            errors.append(
                f"line {line_of(source, row.start())}: row {row_number} uses {len(active)} active hue families "
                f"({', '.join(sorted(active))}); maximum is 4 plus gray"
            )

    # Declared dimensions are intentionally allowed to be illustrative. This
    # lookup merely keeps the declaration parse exercised and future-proof.
    del dimensions
    return errors


def main() -> int:
    parser = argparse.ArgumentParser(description="lint a supertensor .tex source")
    parser.add_argument("sources", nargs="+", type=Path)
    args = parser.parse_args()

    failed = False
    for path in args.sources:
        if not path.is_file():
            print(f"supertensor-lint: no such file: {path}", file=sys.stderr)
            failed = True
            continue
        errors = lint(path)
        if errors:
            failed = True
            print(f"!! {path}", file=sys.stderr)
            for error in errors:
                print(f"   {error}", file=sys.stderr)
        else:
            print(f"  lint  {path}")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())