pigway-frame-flow/pigway_frame_flow/extensions/layout.py

"""Bounded layout plans shared by deterministic and optional AI renderers."""

import json
from io import BytesIO
from typing import Literal

from PIL import Image, ImageDraw, ImageOps
from pydantic import Field, model_validator

from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.models import StrictModel
from pigway_frame_flow.settings import secret


class Region(StrictModel):
    asset_id: str = Field(pattern=r"^[a-f0-9]{32}$")
    x: int = Field(ge=0)
    y: int = Field(ge=0)
    width: int = Field(ge=100)
    height: int = Field(ge=100)


class LayoutInput(StrictModel):
    asset_ids: list[str] = Field(min_length=1, max_length=12)
    title: str = Field(default="", max_length=100)
    message: str = Field(default="", max_length=500)
    orientation: Literal["portrait", "landscape"] = "portrait"
    style: Literal["grid", "hero", "poster"] = "grid"
    use_ai: bool = False
    prompt: str = Field(default="", max_length=1500)

    @model_validator(mode="after")
    def ids(self):
        import re

        if len(set(self.asset_ids)) != len(self.asset_ids) or any(
            not re.fullmatch(r"[a-f0-9]{32}", x) for x in self.asset_ids
        ):
            raise ValueError("Choose distinct local images")
        return self


class LayoutPlan(StrictModel):
    width: int = Field(ge=1200, le=1600)
    height: int = Field(ge=1200, le=1600)
    title: str = Field(default="", max_length=100)
    message: str = Field(default="", max_length=500)
    regions: list[Region] = Field(min_length=1, max_length=12)
    engine: Literal["local", "ai"] = "local"

    @model_validator(mode="after")
    def bounds(self):
        if (self.width, self.height) not in [(1200, 1600), (1600, 1200)]:
            raise ValueError("Unsupported canvas")
        for i, region in enumerate(self.regions):
            if (
                region.x + region.width > self.width
                or region.y + region.height > self.height - 120
                or region.y < 150
            ):
                raise ValueError("Photo region outside safe photo area")
            for other in self.regions[:i]:
                if (
                    region.x < other.x + other.width
                    and other.x < region.x + region.width
                    and region.y < other.y + other.height
                    and other.y < region.y + region.height
                ):
                    raise ValueError("Photo regions overlap")
        return self


def local_plan(request: LayoutInput):
    import math

    width, height = (1200, 1600) if request.orientation == "portrait" else (1600, 1200)
    gap, top, bottom = 24, 170, height - 140
    n = len(request.asset_ids)
    regions = []
    if request.style in ("hero", "poster") and n > 1:
        hero_height = int((bottom - top) * (0.45 if n > 7 else 0.62))
        regions.append(
            Region(
                asset_id=request.asset_ids[0],
                x=gap,
                y=top,
                width=width - 2 * gap,
                height=hero_height - gap,
            )
        )
        top += hero_height
        ids = request.asset_ids[1:]
    else:
        ids = request.asset_ids
    columns = min(3, max(1, math.ceil(math.sqrt(len(ids) * width / (bottom - top)))))
    rows = math.ceil(len(ids) / columns)
    cell_w, cell_h = (
        (width - gap * (columns + 1)) // columns,
        (bottom - top - gap * (rows - 1)) // rows,
    )
    for i, asset_id in enumerate(ids):
        regions.append(
            Region(
                asset_id=asset_id,
                x=gap + (i % columns) * (cell_w + gap),
                y=top + (i // columns) * (cell_h + gap),
                width=cell_w,
                height=cell_h,
            )
        )
    return LayoutPlan(
        width=width,
        height=height,
        title=request.title,
        message=request.message,
        regions=regions,
    )


class LayoutStudio:
    def __init__(self, renderer):
        self.renderer = renderer

    def plan(self, request):
        plan = local_plan(request)
        metadata = []
        for asset_id in request.asset_ids:
            with Image.open(
                self.renderer.assets.path(asset_id), formats=("JPEG", "PNG", "WEBP")
            ) as image:
                metadata.append(
                    {"asset_id": asset_id, "width": image.width, "height": image.height}
                )
        if not request.use_ai:
            return plan
        db = self.renderer.db
        row = db.one("SELECT payload FROM app_config WHERE key='ai_layout'")
        config = json.loads(row["payload"]) if row else {}
        key = secret("FRAME_FLOW_AI_KEY")
        key_path = self.renderer.settings.config_dir / "ai-layout.secret"
        if not key and key_path.exists():
            key = key_path.read_text().strip()
        if not key or not config.get("endpoint") or not config.get("model"):
            raise FlowError(
                E.CONFIG_REQUIRED,
                "Configure the AI layout endpoint, model and key first",
                409,
            )
        self.renderer.settings.additional_secrets.append(key)
        body = {
            "model": config["model"],
            "response_format": {"type": "json_object"},
            "messages": [
                {
                    "role": "system",
                    "content": "Return ONLY a JSON layout following this example. Keep canvas size, all asset ids exactly once, non-overlapping photo rectangles y>=150 and bottom<=height-120. Keep title/message unchanged. Never include code, URLs or HTML. Example: "
                    + plan.model_dump_json(),
                },
                {
                    "role": "user",
                    "content": json.dumps(
                        {
                            "images": metadata,
                            "style": request.style,
                            "instruction": request.prompt,
                        },
                        ensure_ascii=False,
                    ),
                },
            ],
            "max_tokens": 2000,
        }
        try:
            data = self.renderer.assets.fetch_url(
                config["endpoint"],
                method="POST",
                headers={
                    "Authorization": "Bearer " + key,
                    "Content-Type": "application/json",
                },
                body=json.dumps(body).encode(),
                limit=256 * 1024,
            )
            result = json.loads(data)["choices"][0]["message"]["content"]
            candidate = LayoutPlan.model_validate_json(result)
            if (
                sorted(x.asset_id for x in candidate.regions)
                != sorted(request.asset_ids)
                or (candidate.width, candidate.height) != (plan.width, plan.height)
                or candidate.title != request.title
                or candidate.message != request.message
            ):
                raise ValueError()
            candidate.engine = "ai"
            return candidate
        except Exception:
            raise FlowError(
                E.RENDER_FAILED,
                "AI layout failed or returned an invalid plan; choose local layout to continue",
                502,
            ) from None

    def render(self, data, profile):
        plan = LayoutPlan.model_validate(data["layout"])
        canvas = Image.new("RGB", (plan.width, plan.height), "white")
        for region in plan.regions:
            with Image.open(
                self.renderer.assets.path(region.asset_id),
                formats=("JPEG", "PNG", "WEBP"),
            ) as source:
                image = ImageOps.exif_transpose(source).convert("RGB")
                image = ImageOps.fit(
                    image,
                    (region.width, region.height),
                    method=Image.Resampling.LANCZOS,
                )
                canvas.paste(image, (region.x, region.y))
        draw = ImageDraw.Draw(canvas)
        for text, y, size in [
            (plan.title, 35, 54),
            (plan.message, plan.height - 100, 30),
        ]:
            font = self.renderer.font(size)
            while text and draw.textlength(text, font=font) > plan.width - 48:
                text = text[:-1]
            draw.text((24, y), text, font=font, fill="black")
        if profile == "grayscale":
            canvas = ImageOps.grayscale(canvas).convert("RGB")
        out = BytesIO()
        canvas.save(out, "JPEG", quality=93, optimize=True)
        return out.getvalue()