Blame for core/crello.py

"""Crello row -> ground-truth document + the assets a model gets to work with.

What survives the conversion, and why:

- Image, SVG and SVG-mask layers become `image` elements over pre-rendered
  assets. Their colours are baked into the PNG, exactly as a stock asset would
  arrive in a real editor.
- A `ColoredBackground` that is one flat colour becomes the document
  background; a textured or gradient one stays an asset.
- Text layers become real `text` elements. Crello's pre-rendered text PNGs are
  withheld from the assets, otherwise the cheapest policy is to paste pictures
  of the words instead of typing them.
- Per-character colour, bold and italic become runs. Crello's explicit
  `text_line` indices become `\\n` breaks, and the raw newlines it says to
  ignore are dropped.
- Identical assets are merged, so a design that uses one divider twice asks for
  the same asset to be placed twice rather than rewarding a coin-flip over two
  indistinguishable ids.
"""

from __future__ import annotations

import hashlib
import io
import random
from collections import Counter
from dataclasses import dataclass, field

import numpy as np
from PIL import Image

from .colors import parse_color
from .doc import DesignDoc, Element, Run
from .fonts import FontRegistry
from .text import fit_box

_QUANTISED = ("x", "y", "width", "height", "angle", "opacity", "font_size", "line_height",
              "letter_spacing")

ASSET_KIND = {
    "ImageElement": "photo",
    "SvgMaskElement": "photo",
    "SvgElement": "graphic",
    "ColoredBackground": "background",
}


@dataclass
class ConvertedDesign:
    crello_id: str
    doc: DesignDoc
    assets: dict[str, Image.Image]
    asset_meta: dict[str, dict]
    meta: dict
    preview: Image.Image | None
    warnings: list[str] = field(default_factory=list)


def _label(features, key: str, value):
    """ClassLabel int (or list of ints) -> name, tolerant of already-decoded rows."""
    feat = features[key]
    inner = getattr(feat, "feature", feat)
    names = getattr(inner, "names", None)
    if names is None:
        return value
    if isinstance(value, list):
        return [names[v] if isinstance(v, int) else v for v in value]
    return names[value] if isinstance(value, int) else value


def _to_pil(value) -> Image.Image | None:
    if value is None:
        return None
    if isinstance(value, Image.Image):
        return value
    if isinstance(value, dict) and value.get("bytes"):
        return Image.open(io.BytesIO(value["bytes"]))
    return None


def _flat_color(img: Image.Image) -> str | None:
    arr = np.asarray(img.convert("RGBA").resize((32, 32)), dtype=np.float32)
    opaque = arr[:, :, 3] > 250
    if opaque.mean() < 0.98:
        return None
    rgb = arr[:, :, :3][opaque]
    if rgb.std(axis=0).max() > 2.5:
        return None
    r, g, b = (int(round(v)) for v in rgb.mean(axis=0))
    return f"#{r:02x}{g:02x}{b:02x}"


def _text_runs(text: str, colors, bolds, italics, lines) -> tuple[str, list[dict]]:
    """Rebuild text with explicit line breaks and per-character style runs."""

    def at(seq, i, default):
        try:
            return seq[i]
        except (IndexError, TypeError):
            return default

    out_chars: list[tuple[str, str, bool, bool]] = []
    prev_line = None
    for i, ch in enumerate(text):
        if ch == "\n":
            continue
        line = at(lines, i, prev_line or 0)
        style = (
            parse_color(at(colors, i, "#000000") or "#000000"),
            bool(at(bolds, i, False)),
            bool(at(italics, i, False)),
        )
        if prev_line is not None and line != prev_line:
            out_chars.append(("\n", *out_chars[-1][1:]) if out_chars else ("\n", *style))
        out_chars.append((ch, *style))
        prev_line = line
    runs: list[dict] = []
    for ch, color, bold, italic in out_chars:
        if runs and (runs[-1]["color"], runs[-1]["bold"], runs[-1]["italic"]) == (
            color,
            bold,
            italic,
        ):
            runs[-1]["text"] += ch
        else:
            runs.append({"text": ch, "color": color, "bold": bold, "italic": italic})
    return "".join(r["text"] for r in runs), runs


def convert_row(row: dict, features, fonts: FontRegistry) -> ConvertedDesign:
    crello_id = str(row["id"])
    width, height = int(row["canvas_width"]), int(row["canvas_height"])
    types = _label(features, "type", list(row["type"]))
    font_names = _label(features, "font", list(row["font"]))
    aligns = _label(features, "text_align", list(row["text_align"]))
    doc = DesignDoc(width=width, height=height, background="#ffffff")
    warnings: list[str] = []

    raw_assets: list[tuple[str, Image.Image, str]] = []  # (element slot, image, kind)
    pending: list[tuple[int, dict]] = []  # element kwargs, with asset slot names
    bg_set = False

    for i in range(int(row["length"])):
        etype = types[i]
        x, y = float(row["left"][i]), float(row["top"][i])
        w, h = float(row["width"][i]), float(row["height"][i])
        angle = float(row["angle"][i] or 0.0)
        opacity = float(row["opacity"][i]) if row["opacity"][i] is not None else 1.0
        if w < 1 or h < 1 or opacity <= 0.01:
            warnings.append(f"dropped degenerate element {i} ({etype})")
            continue
        common = dict(x=x, y=y, width=w, height=h, angle=angle, opacity=min(1.0, opacity))

        if etype == "TextElement":
            text = str(row["text"][i] or "")
            if not text.strip():
                continue
            family_raw = font_names[i] or ""
            if not fonts.knows(family_raw):
                warnings.append(f"font fallback: {family_raw!r}")
            family = fonts.resolve(family_raw)
            joined, runs = _text_runs(
                text,
                row["text_color"][i],
                row["font_bold"][i],
                row["font_italic"][i],
                row["text_line"][i],
            )
            counts = Counter((r["color"], r["bold"], r["italic"]) for r in runs for _ in r["text"])
            (color, bold, italic), _ = counts.most_common(1)[0]
            uniform = len({(r["color"], r["bold"], r["italic"]) for r in runs}) == 1
            el_runs = None
            if not uniform:
                el_runs = [
                    Run(
                        text=r["text"],
                        color=None if r["color"] == color else r["color"],
                        bold=None if r["bold"] == bold else r["bold"],
                        italic=None if r["italic"] == italic else r["italic"],
                    )
                    for r in runs
                ]
            align = aligns[i] if aligns[i] in ("left", "center", "right") else "left"
            kwargs = dict(
                kind="text",
                text=joined,
                font_family=family,
                font_size=float(row["font_size"][i] or 24.0),
                color=color,
                bold=bold,
                italic=italic,
                align=align,
                line_height=float(row["line_height"][i] or 1.0),
                letter_spacing=float(row["letter_spacing"][i] or 0.0),
                uppercase=bool(row["capitalize"][i]),
                runs=el_runs,
                **common,
            )
            pending.append((i, kwargs))
            continue

        img = _to_pil(row["image"][i])
        if img is None:
            warnings.append(f"element {i} ({etype}) has no image")
            continue
        if etype == "ColoredBackground" and not bg_set:
            flat = _flat_color(img)
            colors = row["color"][i] or []
            if flat is not None:
                doc.background = parse_color(colors[0]) if colors else flat
                bg_set = True
                continue
        slot = f"s{i}"
        raw_assets.append((slot, img.convert("RGBA"), ASSET_KIND.get(etype, "graphic")))
        pending.append((i, dict(kind="image", asset_id=slot, **common)))

    # Merge identical assets, then hand out ids in a shuffled order so the id
    # sequence does not leak the stacking order of the ground truth.
    digest_to_slot: dict[str, str] = {}
    slot_alias: dict[str, str] = {}
    unique: list[tuple[str, Image.Image, str]] = []
    for slot, img, kind in raw_assets:
        digest = hashlib.sha1(img.tobytes() + str(img.size).encode()).hexdigest()
        if digest in digest_to_slot:
            slot_alias[slot] = digest_to_slot[digest]
            continue
        digest_to_slot[digest] = slot
        slot_alias[slot] = slot
        unique.append((slot, img, kind))
    order = list(range(len(unique)))
    random.Random(int(hashlib.sha1(crello_id.encode()).hexdigest()[:8], 16)).shuffle(order)
    slot_to_id = {unique[k][0]: f"a{n + 1}" for n, k in enumerate(order)}

    assets: dict[str, Image.Image] = {}
    asset_meta: dict[str, dict] = {}
    for slot, img, kind in unique:
        aid = slot_to_id[slot]
        assets[aid] = img
        asset_meta[aid] = {
            "id": aid,
            "kind": kind,
            "aspect": round(img.width / img.height, 4),
        }

    for _, kwargs in pending:
        if kwargs["kind"] == "image":
            kwargs["asset_id"] = slot_to_id[slot_alias[kwargs["asset_id"]]]
        # Crello stores float32. Quantise to what a tool call can express, so
        # the expert trajectory reproduces the ground truth exactly.
        for key in _QUANTISED:
            if key in kwargs and isinstance(kwargs[key], float):
                kwargs[key] = round(kwargs[key], 2)
        el = Element(id=doc.new_id(), **kwargs)
        if el.kind == "text":
            _, el.height = fit_box(el, fonts)
        doc.elements.append(el)

    meta = {
        "crello_id": crello_id,
        "title": str(row.get("title") or ""),
        "format": _label(features, "format", row["format"]),
        "category": _label(features, "category", row["category"]),
        "keywords": list(row.get("keywords") or []),
        "industries": _label(features, "industries", list(row.get("industries") or [])),
        "canvas": {"width": width, "height": height},
    }
    return ConvertedDesign(
        crello_id=crello_id,
        doc=doc,
        assets=assets,
        asset_meta=asset_meta,
        meta=meta,
        preview=_to_pil(row.get("preview")),
        warnings=warnings,
    )