#!/usr/bin/env python3
"""Source lint for superderive figures."""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path

ROLE_DECL_RE = re.compile(r"\\sdsetrole\{([^{}]+)\}\{([^{}]+)\}")
ROLE_KEY_RE = re.compile(r"role\s*=\s*([A-Za-z0-9_-]+)")
RAW_TIKZ_RE = re.compile(r"\\(?:draw|fill|path)\b")


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*superderive-lint:\s*(.+)$", source, re.MULTILINE):
        found.update(item.strip() for item in match.group(1).split(","))
    return found


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

    declared = {
        "neutral": "sdGray",
        "keep": "sdTeal",
        "rewrite": "sdOrange",
        "cancel": "sdCoral",
        "intro": "sdViolet",
        "warn": "sdCoral",
    }
    for match in ROLE_DECL_RE.finditer(source):
        name, color = (part.strip() for part in match.groups())
        previous = declared.get(name)
        if previous is not None and previous != color and name != "neutral":
            errors.append(
                f"line {line_of(source, match.start())}: role {name!r} changes "
                f"from {previous!r} to {color!r}"
            )
        declared[name] = color

    used_colors = {"sdGray"}
    for match in ROLE_KEY_RE.finditer(source):
        role = match.group(1)
        color = declared.get(role)
        if color is None:
            errors.append(
                f"line {line_of(source, match.start())}: role {role!r} is used "
                f"before \\sdsetrole"
            )
        else:
            used_colors.add(color)

    active = used_colors - {"sdGray"}
    if len(active) > 4:
        errors.append(
            f"figure uses {len(active)} active hue families "
            f"({', '.join(sorted(active))}); maximum is 4 plus gray"
        )

    if "allow-raw-tikz" not in allowed:
        for match in RAW_TIKZ_RE.finditer(source):
            errors.append(
                f"line {line_of(source, match.start())}: hand-rolled TikZ "
                f"drawing; use \\sdstep/\\sdcancel/\\sdsubst"
            )

    formula_positions = [m.start() for m in re.finditer(r"\\sdtopformula\b", source)]
    row_ends = [m.start() for m in re.finditer(r"\\sdrowend\b", source)]
    if formula_positions and row_ends and formula_positions[-1] < row_ends[-1]:
        errors.append(
            f"line {line_of(source, formula_positions[-1])}: \\sdtopformula "
            f"must follow the last \\sdrowend"
        )

    cols = list(re.finditer(r"\\sdcol\b", source))
    if len(cols) > 1 and "allow-multiple-contrast" not in allowed:
        errors.append(
            f"line {line_of(source, cols[1].start())}: second \\sdcol; "
            f"budget is one illegal/legal pair per figure"
        )
    return errors


def main() -> int:
    parser = argparse.ArgumentParser(description="lint a superderive .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"superderive-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())