rendering.py

"""Existing evaluation renderer, packaged without evaluation/cache dependencies."""
from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence

from PIL import Image, ImageDraw, ImageFont

@dataclass(frozen=True)
class DocumentImageConfig:
    """Exact document rendering parameters used by the vision evaluation."""

    font_path: Path
    font_index: int = 0
    font_family: str = "Roboto"
    font_style: str = "Regular"
    font_size_px: int = 16
    dpi: int = 96
    line_height_px: int | None = None
    width_px: int = 896
    min_height_px: int = 32
    height_step_px: int = 32
    max_height_px: int = 896
    max_lines_per_page: int | None = None

    def __post_init__(self) -> None:
        object.__setattr__(self, "font_path", Path(self.font_path))
        line_height_px = self.font_size_px if self.line_height_px is None else self.line_height_px
        object.__setattr__(self, "line_height_px", line_height_px)
        max_lines_per_page = (
            min(56, self.max_height_px // line_height_px)
            if self.max_lines_per_page is None
            else self.max_lines_per_page
        )
        object.__setattr__(self, "max_lines_per_page", max_lines_per_page)
        numeric_values = {
            "font_index": self.font_index,
            "font_size_px": self.font_size_px,
            "dpi": self.dpi,
            "line_height_px": line_height_px,
            "width_px": self.width_px,
            "min_height_px": self.min_height_px,
            "height_step_px": self.height_step_px,
            "max_height_px": self.max_height_px,
            "max_lines_per_page": max_lines_per_page,
        }
        for name, value in numeric_values.items():
            if name == "font_index" and value < 0:
                raise ValueError("font_index cannot be negative")
            if name != "font_index" and value < 1:
                raise ValueError(f"{name} must be positive")
        if self.min_height_px > self.max_height_px:
            raise ValueError("min_height_px cannot exceed max_height_px")
        if self.min_height_px % self.height_step_px or self.max_height_px % self.height_step_px:
            raise ValueError("page heights must be exact multiples of height_step_px")
        if line_height_px * max_lines_per_page > self.max_height_px:
            raise ValueError("max_lines_per_page does not fit in max_height_px")

    def to_dict(self) -> dict[str, Any]:
        rendering = {
            "font": f"{self.font_family} {self.font_style}",
            "font_path": str(self.font_path),
            "font_size_px": self.font_size_px,
            "dpi": self.dpi,
            "line_height_px": self.line_height_px,
            "width_px": self.width_px,
            "min_height_px": self.min_height_px,
            "height_step_px": self.height_step_px,
            "max_height_px": self.max_height_px,
            "max_lines_per_page": self.max_lines_per_page,
            "margins_px": 0,
            "background": "white",
            "foreground": "black",
            "whitespace_normalization": "all_runs_to_single_space",
            "wrap_measurement": (
                "roboto_pixel_width"
                if self.font_family.casefold() == "roboto"
                else "font_pixel_width"
            ),
            "overwide_word_split": "character",
            "page_overlap": 0,
        }
        if self.font_index:
            rendering["font_index"] = self.font_index
        return rendering


class DocumentImageRenderer:
    """Normalize, pixel-wrap, paginate, and render document text."""

    def __init__(self, config: DocumentImageConfig) -> None:
        self.config = config
        if not config.font_path.is_file():
            raise FileNotFoundError(f"font not found: {config.font_path}")
        self.font = ImageFont.truetype(
            str(config.font_path),
            config.font_size_px,
            index=config.font_index,
        )
        family, style = self.font.getname()
        if (
            family.casefold() != config.font_family.casefold()
            or style.casefold() != config.font_style.casefold()
        ):
            raise ValueError(
                f"{config.font_family} {config.font_style} is required, "
                f"but {config.font_path} reports {family} {style}"
            )

    @staticmethod
    def normalize_text(text: str) -> str:
        return " ".join(str(text).split())

    def text_length(self, text: str) -> float:
        return self.font.getlength(text)

    def draw_line(self, draw: ImageDraw.ImageDraw, line: str, y: int) -> None:
        draw.text(
            (0, y),
            line,
            font=self.font,
            fill="black",
            anchor="lt",
        )

    def rendering_metadata(self) -> dict[str, Any]:
        return self.config.to_dict()

    def _split_overwide_word(self, word: str, max_width: int) -> list[str]:
        chunks: list[str] = []
        offset = 0
        while offset < len(word):
            low, high = offset + 1, len(word)
            best = offset
            while low <= high:
                middle = (low + high) // 2
                if self.text_length(word[offset:middle]) <= max_width:
                    best = middle
                    low = middle + 1
                else:
                    high = middle - 1
            if best == offset:
                raise ValueError(
                    f"width {max_width}px is narrower than one rendered character"
                )
            chunks.append(word[offset:best])
            offset = best
        return chunks

    def wrap_text(self, text: str, *, max_width: int | None = None) -> list[str]:
        max_width = self.config.width_px if max_width is None else max_width
        if max_width < 1:
            raise ValueError("max_width must be positive")
        normalized = self.normalize_text(text)
        if not normalized:
            return []

        lines: list[str] = []
        current = ""
        for word in normalized.split(" "):
            if self.text_length(word) > max_width:
                if current:
                    lines.append(current)
                    current = ""
                chunks = self._split_overwide_word(word, max_width)
                lines.extend(chunks[:-1])
                current = chunks[-1]
                continue

            candidate = word if not current else f"{current} {word}"
            if self.text_length(candidate) <= max_width:
                current = candidate
            else:
                lines.append(current)
                current = word
        if current:
            lines.append(current)
        return lines

    def paginate_lines(self, lines: Sequence[str]) -> list[list[str]]:
        if not lines:
            return [[]]
        page_size = self.config.max_lines_per_page
        return [list(lines[start : start + page_size]) for start in range(0, len(lines), page_size)]

    def page_height(self, line_count: int) -> int:
        if line_count < 0 or line_count > self.config.max_lines_per_page:
            raise ValueError(
                f"line_count must be between 0 and {self.config.max_lines_per_page}"
            )
        content_height = max(1, line_count) * self.config.line_height_px
        stepped_height = math.ceil(content_height / self.config.height_step_px) * self.config.height_step_px
        return min(
            self.config.max_height_px,
            max(self.config.min_height_px, stepped_height),
        )

    def render_document(self, text: str) -> list[Image.Image]:
        lines = self.wrap_text(text)
        rendered: list[Image.Image] = []
        for page_lines in self.paginate_lines(lines):
            image = Image.new(
                "RGB",
                (self.config.width_px, self.page_height(len(page_lines))),
                "white",
            )
            draw = ImageDraw.Draw(image)
            for line_index, line in enumerate(page_lines):
                self.draw_line(draw, line, line_index * self.config.line_height_px)
            rendered.append(image)
        return rendered



def prepare_render_inputs(inputs: dict, renderer: DocumentImageRenderer) -> dict:
    """Convert render documents to images; leave query/instruction/scoring untouched."""
    if not isinstance(inputs, dict):
        raise TypeError("inputs must be a dictionary")
    documents = inputs.get("documents", [])
    if not isinstance(documents, list):
        raise TypeError("documents must be a list")
    converted = []
    for document in documents:
        if not isinstance(document, dict):
            raise TypeError("each document must be a dictionary")
        if "render" not in document:
            converted.append(document)
            continue
        if any(key in document for key in ("text", "image", "video")):
            raise ValueError("render cannot be combined with text, image, or video")
        if not isinstance(document["render"], str):
            raise TypeError("render must contain document text as a string")
        converted.append({
            **{key: value for key, value in document.items() if key != "render"},
            "image": renderer.render_document(document["render"]),
        })
    return {**inputs, "documents": converted}