scripts/render_visual_tables.py

"""Render portable, fail-closed representative cross-codec visual tables."""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import logging
import os
import stat
import textwrap
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw, ImageFont

LOGGER = logging.getLogger(__name__)
CODECS: dict[str, tuple[str, str, str, str]] = {
    "jpeg_q100": ("JPEG Q100", "jpeg_q100_raw", "jpeg_q100_repair_zero", "jpeg_q100_protected_high"),
    "jpeg_q80": ("JPEG Q80", "jpeg_q80_raw", "jpeg_q80_repair_zero", "jpeg_q80_protected_high"),
    "webp_q80": ("WebP Q80", "webp_q80_raw", "webp_q80_repair_zero", "webp_q80_protected_high"),
    "avif_q70": ("AVIF Q70", "avif_q70_raw", "avif_q70_repair_zero", "avif_q70_protected_high"),
    "jxl_d1": ("JPEG XL d1", "jxl_d1_raw", "jxl_d1_repair_zero", "jxl_d1_protected_high"),
}
COLUMNS = (("Baseline PNG", "png_baseline"), ("WebP Lossless", "webp_lossless"))
BG, PANEL, TEXT, MUTED, ACCENT, FAIL = "#0d141b", "#111b24", "#edf5fb", "#9fb4c3", "#8bd5f7", "#ffb1c6"


def _font(size: int, mono: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    candidates = (
        ("/System/Library/Fonts/SFNSMono.ttf", "/System/Library/Fonts/SFNS.ttf")
        if mono
        else ("/System/Library/Fonts/SFNS.ttf",)
    )
    for candidate in candidates:
        if Path(candidate).is_file():
            return ImageFont.truetype(candidate, size)
    return ImageFont.load_default()


def load_prompt_ids(path: str | Path) -> list[str]:
    """Load an ordered prompt-id manifest without guessing missing prompts."""
    value = json.loads(Path(path).read_text(encoding="utf-8"))
    ids = value.get("prompt_ids") if isinstance(value, dict) else value
    if not isinstance(ids, list) or not ids or not all(isinstance(item, (str, int)) for item in ids):
        raise ValueError("prompt-id manifest must contain a non-empty prompt_ids list")
    result = [str(item) for item in ids]
    if len(result) != len(set(result)):
        raise ValueError("prompt-id manifest contains duplicates")
    return result


def _sample_hash(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def _read_sample_bytes_no_follow(root: Path, relative: str) -> bytes:
    """Open every path component without following symlinks and return immutable bytes."""
    candidate = Path(relative)
    if candidate.is_absolute() or not candidate.parts or any(part in {"", ".", ".."} for part in candidate.parts):
        raise ValueError(f"path escapes run: {relative}")
    required_flags = ("O_DIRECTORY", "O_NOFOLLOW", "O_NONBLOCK")
    if any(not hasattr(os, name) for name in required_flags):
        raise ValueError("safe sample opening is unsupported on this platform")
    no_follow = os.O_NOFOLLOW
    directory_flags = os.O_RDONLY | os.O_DIRECTORY | no_follow
    directory_fd = -1
    file_fd = -1
    try:
        directory_fd = os.open(root, directory_flags)
        for part in candidate.parts[:-1]:
            next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
            os.close(directory_fd)
            directory_fd = next_fd
        file_fd = os.open(candidate.parts[-1], os.O_RDONLY | no_follow | os.O_NONBLOCK, dir_fd=directory_fd)
        if not stat.S_ISREG(os.fstat(file_fd).st_mode):
            raise ValueError(f"successful sample is not a regular file: {relative}")
        with os.fdopen(file_fd, "rb") as handle:
            file_fd = -1
            return handle.read()
    except OSError as exc:
        raise ValueError(f"successful sample cannot be opened without following symlinks: {relative}") from exc
    finally:
        if file_fd >= 0:
            os.close(file_fd)
        if directory_fd >= 0:
            os.close(directory_fd)


def _safe_run_path(root: Path, relative: str) -> Path:
    candidate = Path(relative)
    if candidate.is_absolute():
        raise ValueError(f"path escapes run: {relative}")
    raw_path = root / candidate
    if raw_path.is_symlink():
        raise ValueError(f"path is a symlink: {relative}")
    path = raw_path.resolve()
    try:
        path.relative_to(root)
    except ValueError as exc:
        raise ValueError(f"path escapes run: {relative}") from exc
    return path


def load_rows(run_dir: str | Path, prompt_ids: list[str]) -> dict[tuple[str, str], dict[str, Any]]:
    """Load a complete generation matrix and reject unverifiable success rows."""
    root = Path(run_dir).resolve()
    path = _safe_run_path(root, "results/raw_rows.jsonl")
    if path.is_symlink() or not path.is_file():
        raise ValueError("raw_rows.jsonl is missing or is a symlink")
    rows: dict[tuple[str, str], dict[str, Any]] = {}
    conditions = {"png_baseline", "webp_lossless"} | {item for codec in CODECS.values() for item in codec[1:]}
    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
        if not line.strip():
            continue
        row = json.loads(line)
        if row.get("row_type") != "generation":
            continue
        condition, prompt_id = str(row.get("condition")), str(row.get("prompt_id"))
        if condition not in conditions or prompt_id not in prompt_ids:
            continue
        key = (condition, prompt_id)
        if key in rows:
            raise ValueError(f"duplicate generation row: {condition}/{prompt_id}")
        if row.get("success"):
            sample = row.get("sample_path")
            if not isinstance(sample, str) or not sample:
                raise ValueError(f"successful row has no sample path: {condition}/{prompt_id}")
            sample_path = _safe_run_path(root, sample)
            if not sample_path.is_file() or sample_path.is_symlink():
                raise ValueError(f"successful sample is missing: {sample}")
            recorded = row.get("sample_sha256", row.get("sample_hash"))
            if recorded is not None and recorded != _sample_hash(sample_path):
                raise ValueError(f"sample hash mismatch: {sample}")
        else:
            if row.get("sample_path") or row.get("sample_sha256") or row.get("sample_hash"):
                raise ValueError(f"failed row contains sample provenance: {condition}/{prompt_id}")
        rows[key] = row
    missing = [
        (condition, prompt_id)
        for condition in sorted(conditions)
        for prompt_id in prompt_ids
        if (condition, prompt_id) not in rows
    ]
    if missing:
        raise ValueError(f"generation matrix incomplete; first missing cell: {missing[0]}")
    return rows


def _failure_label(row: dict[str, Any]) -> str:
    if row.get("blocked"):
        return "NO SAMPLE\n\nWEIGHT GATE\nnon-finite weights"
    error = str(row.get("error") or row.get("first_error") or "generation failed")
    return "NO SAMPLE\n\nGENERATION\n" + ("non-finite output" if "non-finite" in error.lower() else "generation failed")


def _draw_cell(
    canvas: Image.Image, draw: ImageDraw.ImageDraw, root: Path, row: dict[str, Any], x: int, y: int, size: int
) -> None:
    if row.get("success"):
        sample_bytes = _read_sample_bytes_no_follow(root, str(row["sample_path"]))
        recorded_bytes = row.get("sample_bytes")
        if recorded_bytes is not None and (not isinstance(recorded_bytes, int) or recorded_bytes != len(sample_bytes)):
            raise ValueError(f"sample byte count mismatch: {row['sample_path']}")
        recorded_hash = row.get("sample_sha256", row.get("sample_hash"))
        if recorded_hash is not None and recorded_hash != hashlib.sha256(sample_bytes).hexdigest():
            raise ValueError(f"sample hash mismatch: {row['sample_path']}")
        with Image.open(io.BytesIO(sample_bytes)) as image:
            canvas.paste(image.convert("RGB").resize((size, size), Image.Resampling.LANCZOS), (x, y))
        draw.rectangle((x, y, x + size, y + size), outline="#263746")
        if row.get("clip_score") is not None:
            draw.text((x, y + size + 6), f"CLIP {float(row['clip_score']):.3f}", font=_font(14, True), fill=ACCENT)
    else:
        draw.rectangle((x, y, x + size, y + size), fill="#351d28", outline="#e96a91", width=2)
        draw.multiline_text((x + 12, y + 12), _failure_label(row), font=_font(14, True), fill=FAIL, spacing=4)


def render(
    codec_key: str,
    codec: tuple[str, str, str, str],
    root: Path,
    prompts: list[dict[str, Any]],
    rows: dict[tuple[str, str], dict[str, Any]],
    output: Path,
) -> Path:
    """Render one codec comparison table; failed cells remain explicit panels."""
    columns = COLUMNS + (
        (f"{codec[0]} Raw", codec[1]),
        ("Finite Repair Zero", codec[2]),
        ("Protected High Byte", codec[3]),
    )
    size, label_width, gap, row_height = 220, 285, 30, 292
    canvas = Image.new(
        "RGB", (24 + label_width + len(columns) * (size + gap), 160 + len(prompts) * row_height + 70), BG
    )
    draw = ImageDraw.Draw(canvas)
    draw.text((24, 20), f"PixelModel-v4 · {codec[0]} robustness", font=_font(25), fill=TEXT)
    draw.text((24, 57), "100 prompts · 50 steps · CFG 6.0 · representative prompts", font=_font(15, True), fill=MUTED)
    top = 150
    for index, (title, _) in enumerate(columns):
        x = 24 + label_width + index * (size + gap)
        draw.multiline_text((x, 95), title, font=_font(16), fill=ACCENT, spacing=3)
    for index, prompt in enumerate(prompts):
        y = top + index * row_height
        draw.text((24, y + 5), str(prompt.get("category", "")).upper(), font=_font(14), fill=ACCENT)
        draw.text((24, y + 30), str(prompt["prompt_id"]), font=_font(13, True), fill=TEXT)
        wrapped = "\n".join(textwrap.wrap(str(prompt.get("prompt", "")), width=34))
        draw.multiline_text((24, y + 58), wrapped, font=_font(13), fill=TEXT, spacing=3)
        prompt_id = str(prompt["prompt_id"])
        for column, (_, condition) in enumerate(columns):
            _draw_cell(
                canvas, draw, root, rows[(condition, prompt_id)], 24 + label_width + column * (size + gap), y, size
            )
    destination = output / f"visual_table_{codec_key}.png"
    destination.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(destination, format="PNG", optimize=True)
    return destination


def render_contact_sheet(tables: list[Path], output: Path, *, gap: int = 48) -> Path:
    """Stack full-width tables vertically so labels remain readable in browser viewers."""
    if not tables:
        raise ValueError("at least one table is required")
    images: list[Image.Image] = []
    for table in tables:
        with Image.open(table) as image:
            images.append(image.convert("RGB"))
    width = max(image.width for image in images)
    height = sum(image.height for image in images) + gap * (len(images) - 1)
    canvas = Image.new("RGB", (width, height), BG)
    y = 0
    for index, image in enumerate(images):
        canvas.paste(image, ((width - image.width) // 2, y))
        y += image.height
        if index < len(images) - 1:
            divider_y = y + gap // 2
            ImageDraw.Draw(canvas).rectangle((24, divider_y - 1, width - 24, divider_y + 1), fill="#263746")
            y += gap
    destination = output / "visual_tables_contact_sheet.png"
    destination.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(destination, format="PNG", optimize=True)
    return destination


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-dir", type=Path, required=True)
    parser.add_argument("--prompt-ids-manifest", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    prompt_ids = load_prompt_ids(args.prompt_ids_manifest)
    rows = load_rows(args.run_dir, prompt_ids)
    prompts = [rows[("png_baseline", prompt_id)] for prompt_id in prompt_ids]
    output = args.output
    rendered = [render(key, codec, args.run_dir.resolve(), prompts, rows, output) for key, codec in CODECS.items()]
    render_contact_sheet(rendered, output)
    LOGGER.info("rendered %d tables and a vertical contact sheet to %s", len(CODECS), output)
    return 0


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