validate.py

"""Strict parser and Tier-1 mechanical checks for SpriteBench wire-format sprites.

Wire format (fixed for all models):

    PALETTE
    . transparent
    K 34,32,52
    R 214,48,74
    (... max 12 entries, single-char keys)
    GRID 32x32
    <rows of chars, one per pixel>

Token-stable local training/app data may also put a single ASCII space between
cells in grid rows (``K R R K``). Parsed sprites always store compact rows
internally (``KRRK``), so existing render/scoring code sees one representation.

The parser tolerates prose / reasoning / code fences BEFORE the PALETTE line
and AFTER the grid rows, but is STRICT in between. The first violation wins.
Error strings are fed back to models on retry, so they must be specific and
actionable.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:  # pragma: no cover - import only for type checking
    from tasks import Task


# Palette line patterns. Single space separators; trailing whitespace tolerated
# (stripped before matching). Single-char (non-whitespace) key.
_RGB_LINE_RE = re.compile(r"^(\S) (\d{1,3}),(\d{1,3}),(\d{1,3})$")
_TRANSPARENT_LINE_RE = re.compile(r"^(\S) transparent$")
_GRID_LINE_RE = re.compile(r"^GRID (\d+)x(\d+)$")

MAX_PALETTE = 12


def _normalize_grid_row(row: str, width: int, row_num: int) -> tuple[Optional[str], Optional[str]]:
    """Return compact row text, accepting either compact or space-separated cells."""
    if len(row) == width:
        return row, None

    if " " in row:
        parts = row.split(" ")
        if (
            len(parts) == width
            and all(len(part) == 1 for part in parts)
            and all(part != "" for part in parts)
        ):
            return "".join(parts), None
        return None, (
            f"row {row_num} of grid is not a valid compact row or "
            f"space-separated row; expected {width} characters or {width} "
            "single-character cells separated by single spaces"
        )

    return None, (
        f"row {row_num} of grid has {len(row)} characters, "
        f"expected {width}"
    )


@dataclass
class Sprite:
    """A parsed palette-indexed sprite.

    palette maps a single-char key to an (R, G, B) tuple, or to None for the
    transparent entry (whose key must be '.').
    """

    palette: dict[str, Optional[tuple[int, int, int]]]
    width: int
    height: int
    rows: list[str]

    def to_dict(self) -> dict:
        return {
            # JSON has no tuples; store RGB as a list, transparent as null.
            "palette": {
                key: (None if rgb is None else list(rgb))
                for key, rgb in self.palette.items()
            },
            "width": self.width,
            "height": self.height,
            "rows": list(self.rows),
        }

    @staticmethod
    def from_dict(d: dict) -> "Sprite":
        palette: dict[str, Optional[tuple[int, int, int]]] = {}
        for key, rgb in d["palette"].items():
            if rgb is None:
                palette[key] = None
            else:
                palette[key] = (int(rgb[0]), int(rgb[1]), int(rgb[2]))
        return Sprite(
            palette=palette,
            width=int(d["width"]),
            height=int(d["height"]),
            rows=list(d["rows"]),
        )


def parse_sprite(text: str) -> tuple[Optional[Sprite], Optional[str]]:
    """Parse wire-format text into a Sprite.

    Returns (sprite, None) on success or (None, error_message) on failure.
    The error message is specific and actionable for model-retry feedback.

    Tolerance: arbitrary content before the first exact `PALETTE` line and
    after the final grid row is ignored. Everything in between is strict, and
    the first violation determines the returned error.
    """
    if text is None:
        return None, "no PALETTE line found in output"

    lines = text.splitlines()

    # 1. Locate the first line that is exactly `PALETTE` after stripping
    #    surrounding whitespace. Tolerate any prose / fences before it.
    palette_index: Optional[int] = None
    for i, line in enumerate(lines):
        if line.strip() == "PALETTE":
            palette_index = i
            break
    if palette_index is None:
        return None, "no PALETTE line found in output"

    # 2. Parse palette entries on the consecutive lines after PALETTE, up to the
    #    GRID line. Strict from here on.
    palette: dict[str, Optional[tuple[int, int, int]]] = {}
    idx = palette_index + 1
    grid_index: Optional[int] = None

    while idx < len(lines):
        # Strip trailing whitespace only (single-space separators are strict);
        # leading whitespace is part of the strict region and not allowed.
        raw = lines[idx]
        line = raw.rstrip()

        if line.startswith("GRID") or _GRID_LINE_RE.match(line):
            grid_index = idx
            break

        if line == "":
            return None, (
                "blank line inside the palette section is not allowed; "
                "list palette entries then the GRID line with no gaps"
            )

        rgb_match = _RGB_LINE_RE.match(line)
        transparent_match = _TRANSPARENT_LINE_RE.match(line)

        if transparent_match:
            key = transparent_match.group(1)
            if key != ".":
                return None, (
                    f"transparent palette key must be '.' but found '{key}'"
                )
            if key in palette:
                return None, f"duplicate palette key '{key}'"
            palette[key] = None
        elif rgb_match:
            key = rgb_match.group(1)
            r = int(rgb_match.group(2))
            g = int(rgb_match.group(3))
            b = int(rgb_match.group(4))
            for component, value in (("R", r), ("G", g), ("B", b)):
                if value > 255:
                    return None, (
                        f"palette key '{key}' has {component} value {value}, "
                        f"which is out of range (RGB components must be 0-255)"
                    )
            if key in palette:
                return None, f"duplicate palette key '{key}'"
            palette[key] = (r, g, b)
        else:
            return None, (
                f"palette line '{line}' is malformed; expected "
                f"'<key> R,G,B' (e.g. 'K 34,32,52') or '. transparent'"
            )

        if len(palette) > MAX_PALETTE:
            return None, (
                f"palette has {len(palette)} entries, exceeding the maximum "
                f"of {MAX_PALETTE}"
            )

        idx += 1

    if not palette:
        return None, "palette has no entries; declare at least one color"

    # 3. The GRID line must exist and be well formed.
    if grid_index is None:
        return None, "no 'GRID WxH' line found after the palette"

    grid_line = lines[grid_index].rstrip()
    grid_match = _GRID_LINE_RE.match(grid_line)
    if not grid_match:
        return None, (
            f"GRID line '{grid_line}' is malformed; expected 'GRID WxH' "
            f"(e.g. 'GRID 32x32')"
        )

    width = int(grid_match.group(1))
    height = int(grid_match.group(2))
    if width <= 0 or height <= 0:
        return None, (
            f"GRID dimensions must be positive but found {width}x{height}"
        )

    # 4. Exactly H consecutive non-empty rows, each exactly W chars, all chars
    #    declared in the palette. Stop at H rows; trailing content ignored.
    rows: list[str] = []
    row_idx = grid_index + 1
    for r in range(height):
        if row_idx >= len(lines):
            return None, (
                f"grid ended early: expected {height} rows but found "
                f"{len(rows)}"
            )
        # Rows are taken verbatim (no stripping): a trailing space is a real
        # pixel-count error and must be reported, not silently trimmed.
        row = lines[row_idx]
        if row == "":
            return None, (
                f"row {r + 1} of grid is empty; expected {width} characters"
            )
        row, row_err = _normalize_grid_row(row, width, r + 1)
        if row_err is not None:
            return None, row_err
        assert row is not None
        for col, ch in enumerate(row):
            if ch not in palette:
                return None, (
                    f"character '{ch}' at row {r + 1}, column {col + 1} is "
                    f"not declared in the palette"
                )
        rows.append(row)
        row_idx += 1

    sprite = Sprite(palette=palette, width=width, height=height, rows=rows)
    return sprite, None


def footprint(sprite: Sprite) -> set[tuple[int, int]]:
    """Return the set of (x, y) coordinates of non-transparent cells.

    A cell is non-transparent when its palette value is not None. x is the
    column index, y is the row index, both zero-based.
    """
    cells: set[tuple[int, int]] = set()
    for y, row in enumerate(sprite.rows):
        for x, ch in enumerate(row):
            if sprite.palette.get(ch) is not None:
                cells.add((x, y))
    return cells


@dataclass
class Tier1:
    """Tier-1 mechanical check result for a single sample."""

    parse_ok: bool
    dims_ok: bool
    palette_ok: bool
    transparency_ok: bool
    footprint_ok: Optional[bool]  # T10 only, else None
    passed: bool
    errors: list[str]

    def to_dict(self) -> dict:
        return {
            "parse_ok": self.parse_ok,
            "dims_ok": self.dims_ok,
            "palette_ok": self.palette_ok,
            "transparency_ok": self.transparency_ok,
            "footprint_ok": self.footprint_ok,
            "passed": self.passed,
            "errors": list(self.errors),
        }

    @staticmethod
    def from_dict(d: dict) -> "Tier1":
        return Tier1(
            parse_ok=bool(d["parse_ok"]),
            dims_ok=bool(d["dims_ok"]),
            palette_ok=bool(d["palette_ok"]),
            transparency_ok=bool(d["transparency_ok"]),
            footprint_ok=(
                None if d["footprint_ok"] is None else bool(d["footprint_ok"])
            ),
            passed=bool(d["passed"]),
            errors=list(d["errors"]),
        )


def tier1_check(
    sprite: Optional[Sprite], parse_error: Optional[str], task: "Task"
) -> Tier1:
    """Run the deterministic Tier-1 checks for one sample.

    Checks: parse success, exact dimensions, palette adherence, transparency
    sanity (0 < transparent fraction < 1), and for the T10 edit task, exact
    footprint match against task.expected_footprint.

    parse_error is the message returned by parse_sprite (None on success).
    When parse failed on an undeclared-char error, palette_ok is reported False;
    on any other parse failure palette_ok is also False because the palette
    could not be validated. When parse succeeded, the parser guarantees
    palette adherence so palette_ok is True.
    """
    is_edit = getattr(task, "kind", None) == "edit"

    if sprite is None:
        message = parse_error if parse_error else "parse failed"
        # Distinguish an undeclared-char failure so palette_ok reflects it; in
        # all parse-failure cases palette adherence is unverified -> False.
        return Tier1(
            parse_ok=False,
            dims_ok=False,
            palette_ok=False,
            transparency_ok=False,
            footprint_ok=(False if is_edit else None),
            passed=False,
            errors=[message],
        )

    errors: list[str] = []

    # Dimensions: sprite must be exactly task.size x task.size (square).
    dims_ok = sprite.width == task.size and sprite.height == task.size
    if not dims_ok:
        errors.append(
            f"sprite dimensions are {sprite.width}x{sprite.height}, "
            f"expected {task.size}x{task.size}"
        )

    # Palette adherence is guaranteed by the parser on a successful parse.
    palette_ok = True

    # Transparency sanity: must have SOME transparency and SOME content.
    total_cells = sprite.width * sprite.height
    content_cells = footprint(sprite)
    transparent_count = total_cells - len(content_cells)
    transparent_fraction = (
        transparent_count / total_cells if total_cells else 0.0
    )
    transparency_ok = 0.0 < transparent_fraction < 1.0
    if not transparency_ok:
        if transparent_fraction <= 0.0:
            errors.append(
                "sprite has no transparent cells; a sprite must include some "
                "transparency (use the '.' key for background pixels)"
            )
        else:
            errors.append(
                "sprite is entirely transparent; it must contain some "
                "non-transparent content"
            )

    # Footprint match for the T10 edit task only.
    footprint_ok: Optional[bool]
    if is_edit and task.expected_footprint is not None:
        expected = set(task.expected_footprint)
        actual = content_cells
        footprint_ok = actual == expected
        if not footprint_ok:
            added = sorted(actual - expected)
            removed = sorted(expected - actual)
            detail_parts: list[str] = []
            if added:
                detail_parts.append(
                    f"{len(added)} cell(s) became non-transparent that must "
                    f"stay transparent (e.g. {_format_cells(added)})"
                )
            if removed:
                detail_parts.append(
                    f"{len(removed)} cell(s) became transparent that must "
                    f"stay filled (e.g. {_format_cells(removed)})"
                )
            errors.append(
                "footprint changed: the set of transparent cells must remain "
                "exactly identical to the input; "
                + "; ".join(detail_parts)
            )
    elif is_edit:
        # Edit task without an expected footprint declared: cannot verify.
        footprint_ok = None
    else:
        footprint_ok = None

    # passed = all applicable checks true.
    applicable = [dims_ok, palette_ok, transparency_ok]
    if footprint_ok is not None:
        applicable.append(footprint_ok)
    passed = all(applicable)

    return Tier1(
        parse_ok=True,
        dims_ok=dims_ok,
        palette_ok=palette_ok,
        transparency_ok=transparency_ok,
        footprint_ok=footprint_ok,
        passed=passed,
        errors=errors,
    )


def _format_cells(cells: list[tuple[int, int]], limit: int = 5) -> str:
    """Format a short, readable sample of (x, y) cell coordinates."""
    shown = ", ".join(f"({x},{y})" for x, y in cells[:limit])
    if len(cells) > limit:
        shown += ", ..."
    return shown