pigway-frame-flow/pigway_frame_flow/display/libraries.py

"""Source pictures, source-scoped render plans, and single terminal bindings."""

import hashlib
import json
from io import BytesIO
from pathlib import Path
from uuid import uuid4

from PIL import Image
from pydantic import Field, model_validator

from pigway_frame_flow.core.photo_preferences import flags
from pigway_frame_flow.display.thumbnail_plans import ThumbnailPlan
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.models import PhotoSourceContent
from pigway_frame_flow.renderers.images import MAX_IMAGE_PIXELS, convert
from pigway_frame_flow.renderers.photo_caption import render_captioned_photo
from pigway_frame_flow.storage.database import encode, now_iso


class GenerationPlan(ThumbnailPlan):
    source_id: str = Field(min_length=1, max_length=120)
    caption: bool = False
    profile: str = Field(default="color", pattern="^(color|grayscale)$")

    @model_validator(mode="after")
    def caption_size(self):
        if self.caption and min(self.width, self.height) < 128:
            raise ValueError(
                "Caption output requires both dimensions to be at least 128 pixels"
            )
        return self


class Libraries:
    def __init__(self, service):
        self.service, self.db = service, service.db
        for kind in ("source-images", "renditions"):
            (service.settings.data_dir / kind).mkdir(exist_ok=True)
        for sql in [
            "CREATE TABLE IF NOT EXISTS library_plans(id TEXT PRIMARY KEY,source_id TEXT NOT NULL,payload TEXT NOT NULL,archived INTEGER NOT NULL DEFAULT 0)",
            "CREATE TABLE IF NOT EXISTS source_images(id TEXT PRIMARY KEY,source_id TEXT NOT NULL,item_id TEXT NOT NULL,digest TEXT NOT NULL,caption TEXT NOT NULL,created_at TEXT NOT NULL,UNIQUE(source_id,item_id,digest))",
            "CREATE TABLE IF NOT EXISTS library_versions(asset_id TEXT PRIMARY KEY,source_image_id TEXT NOT NULL,plan_id TEXT NOT NULL,revision TEXT NOT NULL,plan TEXT NOT NULL,created_at TEXT NOT NULL,UNIQUE(source_image_id,plan_id,revision))",
            "CREATE TABLE IF NOT EXISTS terminal_bindings(terminal_id TEXT PRIMARY KEY,source_id TEXT NOT NULL,plan_id TEXT NOT NULL)",
            "CREATE TABLE IF NOT EXISTS terminal_deliveries(terminal_id TEXT NOT NULL,asset_id TEXT NOT NULL,operation_id TEXT PRIMARY KEY,parent_operation_id TEXT,created_at TEXT NOT NULL)",
        ]:
            self.db.execute(sql)
        service.operation_handlers["library-generate"] = self.execute

    def terminals(self):
        result = []
        for terminal_id, driver in self.service.registry.drivers.items():
            cached = getattr(driver, "cached", None)
            result.append(
                {
                    "id": terminal_id,
                    "name": getattr(cached, "name", terminal_id),
                    "kind": "frame",
                    "transport_supported": True,
                    "binding": self.binding(terminal_id),
                }
            )
        result.append(
            {
                "id": "watch",
                "name": "Watch",
                "kind": "watch",
                "width": 410,
                "height": 502,
                "transport_supported": False,
                "binding": self.binding("watch"),
            }
        )
        return result

    def binding(self, terminal_id):
        return self.db.one(
            "SELECT * FROM terminal_bindings WHERE terminal_id=?", (terminal_id,)
        )

    def plan(self, plan_id, active=True):
        row = self.db.one("SELECT * FROM library_plans WHERE id=?", (plan_id,))
        if not row or (active and row["archived"]):
            raise FlowError(E.NOT_FOUND, "Generation plan not found", 404)
        return GenerationPlan.model_validate_json(row["payload"])

    def plans(self, source_id):
        self.service.photo_sources.get(source_id)
        return [
            {**json.loads(r["payload"]), "id": r["id"], "archived": bool(r["archived"])}
            for r in self.db.all(
                "SELECT * FROM library_plans WHERE source_id=? ORDER BY archived,id",
                (source_id,),
            )
        ]

    def save_plan(self, plan, plan_id=None):
        with self.service.lock:
            self.service.photo_sources.get(plan.source_id)
            if plan_id:
                previous = self.plan(plan_id)
                if previous.source_id != plan.source_id:
                    raise FlowError(
                        E.CONFLICT,
                        "A generation plan cannot move between source libraries",
                        409,
                    )
                for row in self.db.all(
                    "SELECT terminal_id FROM terminal_bindings WHERE plan_id=?",
                    (plan_id,),
                ):
                    self.validate_size(row["terminal_id"], plan)
            plan_id = plan_id or uuid4().hex
            self.db.execute(
                "INSERT INTO library_plans VALUES(?,?,?,0) ON CONFLICT(id) DO UPDATE SET payload=excluded.payload",
                (plan_id, plan.source_id, encode(plan.model_dump())),
            )
            return {"id": plan_id, **plan.model_dump(), "archived": False}

    def archive(self, plan_id):
        with self.service.lock:
            self.plan(plan_id)
            if self.db.one(
                "SELECT terminal_id FROM terminal_bindings WHERE plan_id=?", (plan_id,)
            ):
                raise FlowError(
                    E.CONFLICT,
                    "Unbind terminals before deleting this generation plan",
                    409,
                    {"reason": "libraryPlanBound"},
                )
            self.db.execute(
                "UPDATE library_plans SET archived=1 WHERE id=?", (plan_id,)
            )

    def validate_size(self, terminal_id, plan):
        if terminal_id == "watch":
            allowed = {(410, 502)}
        else:
            display = self.service.registry.get(terminal_id).device().display
            allowed = {(display.width, display.height), (display.height, display.width)}
        if (plan.width, plan.height) not in allowed:
            raise FlowError(
                E.INVALID_REQUEST,
                "Generation dimensions must match the terminal screen",
                422,
                {"reason": "librarySizeMismatch"},
            )

    def bind(self, terminal_id, source_id, plan_id):
        with self.service.lock:
            self.service.photo_sources.get(source_id)
            plan = self.plan(plan_id)
            if plan.source_id != source_id:
                raise FlowError(
                    E.CONFLICT, "Generation plan belongs to another source library", 409
                )
            self.validate_size(terminal_id, plan)
            self.db.execute(
                "INSERT INTO terminal_bindings VALUES(?,?,?) ON CONFLICT(terminal_id) DO UPDATE SET source_id=excluded.source_id,plan_id=excluded.plan_id",
                (terminal_id, source_id, plan_id),
            )
            return self.binding(terminal_id)

    def ingest(self, source_id, item_id, data, name, caption):
        digest = hashlib.sha256(
            data + b"\0" + encode(caption or {}).encode()
        ).hexdigest()
        row = self.db.one(
            "SELECT id FROM source_images WHERE source_id=? AND item_id=? AND digest=?",
            (source_id, item_id, digest),
        )
        if row:
            return self.picture(row["id"])
        # Normalize EXIF and JPEG storage without resizing or adding decorations.
        try:
            with Image.open(BytesIO(data)) as image:
                if (
                    image.format not in {"JPEG", "PNG", "WEBP"}
                    or image.width * image.height > MAX_IMAGE_PIXELS
                ):
                    raise ValueError("Unsupported source image")
                size = image.size
                if image.getexif().get(274, 1) in (5, 6, 7, 8):
                    size = (image.height, image.width)
            normalized = convert(data, "cover", size=size, quality=95)
        except (OSError, ValueError, Image.DecompressionBombError):
            raise FlowError(E.INVALID_REQUEST, "Invalid source image") from None
        asset = self.service.assets.put(normalized, name, "source-images")
        self.db.execute(
            "INSERT INTO source_images VALUES(?,?,?,?,?,?)",
            (
                asset["asset_id"],
                source_id,
                item_id,
                digest,
                encode(caption or {}),
                now_iso(),
            ),
        )
        self.db.execute(
            "INSERT INTO asset_sources VALUES(?,?)",
            (
                asset["asset_id"],
                encode(
                    {
                        "type": "library_source",
                        "source_id": source_id,
                        "item_id": item_id,
                        "source_rendition": "preview"
                        if self.service.photo_sources.get(source_id).kind
                        == "synology-photos"
                        else "source",
                    }
                ),
            ),
        )
        return self.picture(asset["asset_id"])

    def import_picture(self, source_id, item_id):
        with self.service.lock:
            data, name, caption = self.service.photo_sources.read(source_id, item_id)
            return self.ingest(source_id, item_id, data, name, caption)

    def picture(self, picture_id):
        row = self.db.one("SELECT * FROM source_images WHERE id=?", (picture_id,))
        if not row:
            raise FlowError(E.NOT_FOUND, "Source picture not found", 404)
        return {
            **self.service.assets.info(picture_id),
            **row,
            "caption": json.loads(row["caption"]),
            **flags(self.db, row["source_id"], row["item_id"]),
            "versions": self.versions(picture_id),
        }

    @staticmethod
    def render_parameters(plan):
        return plan.model_dump(exclude={"name", "source_id"})

    def versions(self, picture_id):
        result = []
        rows = self.db.all(
            "SELECT v.*,p.payload AS current_plan,p.archived,a.name AS source_name "
            "FROM library_versions v JOIN library_plans p ON p.id=v.plan_id "
            "JOIN source_images s ON s.id=v.source_image_id JOIN assets a ON a.id=s.id "
            "WHERE v.source_image_id=? ORDER BY v.created_at DESC,v.asset_id DESC",
            (picture_id,),
        )
        for row in rows:
            current = GenerationPlan.model_validate_json(row["current_plan"])
            snapshot = GenerationPlan.model_validate_json(row["plan"])
            asset = self.service.assets.info(row["asset_id"])
            display_name = (
                Path(Path(row["source_name"]).stem + "-" + current.name).name[:110]
                + ".jpg"
            )
            result.append(
                {
                    **asset,
                    "name": display_name,
                    "stored_name": asset["name"],
                    "source_image_id": row["source_image_id"],
                    "plan_id": row["plan_id"],
                    "plan_name": current.name,
                    "plan_archived": bool(row["archived"]),
                    "matches_current_plan": self.render_parameters(current)
                    == self.render_parameters(snapshot),
                    "revision": row["revision"],
                    "plan": json.loads(row["plan"]),
                    "terminal_uses": self.terminal_uses(
                        row["plan_id"], row["asset_id"]
                    ),
                }
            )
        return result

    def terminal_uses(self, plan_id, asset_id):
        return self.db.all(
            "SELECT b.terminal_id,d.operation_id,COALESCE(o.status,'GENERATED') AS status "
            "FROM terminal_bindings b LEFT JOIN terminal_deliveries d ON d.operation_id=(SELECT operation_id FROM terminal_deliveries latest "
            "WHERE latest.terminal_id=b.terminal_id AND latest.asset_id=? ORDER BY latest.created_at DESC,latest.operation_id DESC LIMIT 1) "
            "LEFT JOIN operations o ON o.id=d.operation_id WHERE b.plan_id=? ORDER BY b.terminal_id",
            (asset_id, plan_id),
        )

    def render(self, data, caption, plan):
        if plan.caption and caption:
            return render_captioned_photo(
                data,
                caption,
                self.service.renderer.font,
                plan.fit,
                plan.profile,
                size=(plan.width, plan.height),
                centering=(plan.center_x, plan.center_y),
                quality=plan.quality,
                respect_fit=True,
            )
        return convert(
            data,
            plan.fit,
            plan.profile,
            size=(plan.width, plan.height),
            centering=(plan.center_x, plan.center_y),
            quality=plan.quality,
        )

    def generate(self, picture_id, plan_id, snapshot=None):
        with self.service.lock:
            picture = self.picture(picture_id)
            plan = (
                GenerationPlan.model_validate(snapshot)
                if snapshot
                else self.plan(plan_id)
            )
            if picture["source_id"] != plan.source_id:
                raise FlowError(
                    E.CONFLICT,
                    "Picture and generation plan must share a source library",
                    409,
                )
            parameters = self.render_parameters(plan)
            revision = hashlib.sha256(encode(parameters).encode()).hexdigest()
            row = self.db.one(
                "SELECT asset_id FROM library_versions WHERE source_image_id=? AND plan_id=? AND revision=?",
                (picture_id, plan_id, revision),
            )
            existing = self.versions(picture_id)
            if row:
                return next(v for v in existing if v["asset_id"] == row["asset_id"])
            # Older revisions hashed the display name too. Reuse matching pixels'
            # settings without rewriting historical snapshots or receipt references.
            for version in existing:
                if (
                    version["plan_id"] == plan_id
                    and self.render_parameters(
                        GenerationPlan.model_validate(version["plan"])
                    )
                    == parameters
                ):
                    return version
            data = self.service.assets.path(picture_id).read_bytes()
            asset = self.service.assets.put(
                self.render(data, picture["caption"], plan),
                picture["name"].removesuffix(".jpg") + "-" + plan.name,
                "renditions",
            )
            self.db.execute(
                "INSERT INTO library_versions VALUES(?,?,?,?,?,?)",
                (
                    asset["asset_id"],
                    picture_id,
                    plan_id,
                    revision,
                    encode(plan.model_dump()),
                    now_iso(),
                ),
            )
            self.db.execute(
                "INSERT INTO asset_sources VALUES(?,?)",
                (
                    asset["asset_id"],
                    encode(
                        {
                            "type": "library_version",
                            "source_image_id": picture_id,
                            "source_id": plan.source_id,
                            "plan_id": plan_id,
                            "revision": revision,
                        }
                    ),
                ),
            )
            return next(
                v
                for v in self.versions(picture_id)
                if v["asset_id"] == asset["asset_id"]
            )

    def validate_asset_binding(self, terminal_id, asset_id):
        if self.db.one("SELECT id FROM source_images WHERE id=?", (asset_id,)):
            raise FlowError(
                E.CONFLICT,
                "Generate a terminal version before sending source pixels",
                409,
                {"reason": "libraryGenerateFirst"},
            )
        version = self.db.one(
            "SELECT v.plan_id,s.source_id FROM library_versions v JOIN source_images s ON s.id=v.source_image_id WHERE v.asset_id=?",
            (asset_id,),
        )
        if version:
            binding = self.binding(terminal_id)
            if (
                not binding
                or binding["source_id"] != version["source_id"]
                or binding["plan_id"] != version["plan_id"]
            ):
                raise FlowError(
                    E.CONFLICT, "Image no longer matches the terminal binding", 409
                )

    def send(
        self,
        terminal_id,
        asset_id,
        options=None,
        client_id="admin",
        parent_operation_id=None,
    ):
        with self.service.lock:
            binding = self.binding(terminal_id)
            version = self.db.one(
                "SELECT v.*,s.source_id FROM library_versions v JOIN source_images s ON s.id=v.source_image_id WHERE v.asset_id=?",
                (asset_id,),
            )
            if (
                not binding
                or not version
                or binding["source_id"] != version["source_id"]
                or binding["plan_id"] != version["plan_id"]
            ):
                raise FlowError(
                    E.CONFLICT,
                    "Image does not match the terminal source and plan binding",
                    409,
                )
            if terminal_id == "watch":
                raise FlowError(
                    E.UNSUPPORTED_CAPABILITY, "Watch transfer is not connected", 422
                )
            request = {
                "device_id": terminal_id,
                "content": {"type": "image", "asset_id": asset_id},
                "options": {
                    **(options or {}),
                    "cleanup_unstarred": False,
                    "generate_thumbnail": False,
                    "thumbnail_plan_id": None,
                },
            }
            from pigway_frame_flow.models import DisplayRequest

            request_model = DisplayRequest.model_validate(request)
            request = request_model.model_dump(mode="json")
            previous = self.db.one(
                "SELECT operation_id FROM terminal_deliveries WHERE terminal_id=? AND asset_id=? "
                + ("AND parent_operation_id=? " if parent_operation_id else "")
                + "ORDER BY created_at DESC,operation_id DESC LIMIT 1",
                (terminal_id, asset_id, parent_operation_id)
                if parent_operation_id
                else (terminal_id, asset_id),
            )
            if previous:
                op = self.db.one(
                    "SELECT id,status,payload FROM operations WHERE id=?",
                    (previous["operation_id"],),
                )
                if op and (
                    parent_operation_id
                    or (
                        op["status"]
                        in ("QUEUED", "RENDERING", "UPLOADING", "SYNCING", "UNKNOWN")
                        and json.loads(op["payload"]) == request
                    )
                ):
                    return op["id"]
            self.service.validate(request_model)
            # Reuse the queue's durable idempotency receipt if a restart occurs
            # between submitting a child operation and linking its delivery row.
            key = (
                "library-output-"
                + hashlib.sha256(
                    encode([parent_operation_id, terminal_id, asset_id]).encode()
                ).hexdigest()
                if parent_operation_id
                else None
            )
            operation = self.service.submit(
                "display",
                request,
                terminal_id,
                client_id,
                "library-" + uuid4().hex,
                key=key,
            )
            self.db.execute(
                "INSERT INTO terminal_deliveries VALUES(?,?,?,?,?)",
                (terminal_id, asset_id, operation, parent_operation_id, now_iso()),
            )
            return operation

    def pipeline(
        self,
        source_id,
        item_id=None,
        options=None,
        deliver=False,
        operation_id=None,
        caller="admin",
    ):
        root_operation = (
            (self.db.one("SELECT * FROM operations WHERE id=?", (operation_id,)) or {})
            if operation_id
            else {}
        )
        previous = (
            json.loads(root_operation["result"] or "{}") if root_operation else {}
        )
        bindings = previous.get("binding_snapshot")
        if bindings is None:
            bindings = self.db.all(
                "SELECT * FROM terminal_bindings WHERE source_id=? ORDER BY terminal_id",
                (source_id,),
            )
            bindings = [
                {**b, "plan_snapshot": self.plan(b["plan_id"]).model_dump()}
                for b in bindings
            ]
        if not bindings:
            raise FlowError(
                E.CONFLICT,
                "Bind at least one terminal to this source library first",
                409,
                {"reason": "libraryBindFirst"},
            )
        if operation_id:
            self.service.update(operation_id, "RENDERING")
        if previous.get("source_image_id"):
            selected = self.picture(previous["source_image_id"])
            if selected["source_id"] != source_id:
                raise FlowError(
                    E.CONFLICT, "Saved selection belongs to another source library", 409
                )
        else:
            selected = self.service.content_asset(
                PhotoSourceContent(
                    type="photo_source", source_id=source_id, item_id=item_id
                ),
                operation_id=operation_id,
                source_only=True,
                source_size=max(
                    (
                        (b["plan_snapshot"]["width"], b["plan_snapshot"]["height"])
                        for b in bindings
                    ),
                    key=lambda size: size[0] * size[1],
                ),
            )
        result = {
            "source_image_id": selected["asset_id"],
            "binding_snapshot": bindings,
            "outputs": [],
            "errors": [],
        }
        if operation_id:
            self.service.assets.record_origin(selected["asset_id"], root_operation)
            self.service.update(operation_id, "RENDERING", result)
        versions = {}
        for binding in bindings:
            terminal_id, plan_id = binding["terminal_id"], binding["plan_id"]
            try:
                if plan_id not in versions:
                    versions[plan_id] = self.generate(
                        selected["asset_id"], plan_id, binding["plan_snapshot"]
                    )
                version = versions[plan_id]
                if operation_id:
                    self.service.assets.record_origin(
                        version["asset_id"], root_operation
                    )
                output = {
                    "terminal_id": terminal_id,
                    "asset_id": version["asset_id"],
                    "plan_id": plan_id,
                    "status": "GENERATED",
                }
                if deliver and terminal_id != "watch":
                    output["operation_id"] = self.send(
                        terminal_id,
                        version["asset_id"],
                        options,
                        "library-terminal" if operation_id else caller,
                        parent_operation_id=operation_id,
                    )
                    output["status"] = "QUEUED"
                result["outputs"].append(output)
                if operation_id:
                    self.service.update(operation_id, "RENDERING", result)
            except (FlowError, OSError) as exc:
                result["errors"].append(
                    {
                        "terminal_id": terminal_id,
                        "plan_id": plan_id,
                        "message": str(exc),
                    }
                )
        if operation_id:
            self.service.update(
                operation_id,
                "COMPLETED" if not result["errors"] else "FAILED",
                result,
                error={
                    "code": E.INVALID_REQUEST,
                    "message": "Some terminal outputs failed",
                    "request_id": root_operation["request_id"],
                    "details": {"reason": "libraryOutputPartial"},
                }
                if result["errors"]
                else None,
            )
        return result

    def execute(self, operation, payload):
        result = self.pipeline(
            payload["content"]["source_id"],
            payload["content"].get("item_id"),
            payload.get("options"),
            deliver=True,
            operation_id=operation["id"],
        )
        self.db.history(operation, "COMPLETED" if not result["errors"] else "FAILED")

    def check_asset_references(self, asset_id):
        reference = "%" + asset_id + "%"
        if (
            self.db.one(
                "SELECT id FROM operations WHERE (payload LIKE ? OR result LIKE ?) AND status IN ('QUEUED','RENDERING','UPLOADING','SYNCING','UNKNOWN')",
                (reference, reference),
            )
            or self.db.one("SELECT id FROM jobs WHERE payload LIKE ?", (reference,))
            or self.db.one(
                "SELECT template_id FROM template_config WHERE payload LIKE ?",
                (reference,),
            )
        ):
            raise FlowError(
                E.CONFLICT,
                "Image is used by a task, template or pending operation",
                409,
            )

    def remove_asset(self, asset_id):
        self.service.assets.path(asset_id).unlink(missing_ok=True)
        self.db.execute("DELETE FROM library_versions WHERE asset_id=?", (asset_id,))
        self.db.execute("DELETE FROM terminal_deliveries WHERE asset_id=?", (asset_id,))
        self.db.execute("DELETE FROM source_images WHERE id=?", (asset_id,))
        for table, key in [
            ("asset_sources", "asset_id"),
            ("asset_metadata", "asset_id"),
            ("assets", "id"),
        ]:
            self.db.execute(f"DELETE FROM {table} WHERE {key}=?", (asset_id,))
        self.db.execute(
            "DELETE FROM photo_preferences WHERE source_id='@assets' AND item_id=?",
            (asset_id,),
        )

    def delete_picture(self, picture_id):
        with self.service.lock:
            self.picture(picture_id)
            ids = [
                r["asset_id"]
                for r in self.db.all(
                    "SELECT asset_id FROM library_versions WHERE source_image_id=?",
                    (picture_id,),
                )
            ] + [picture_id]
            for asset_id in ids:
                self.check_asset_references(asset_id)
            for asset_id in ids:
                self.remove_asset(asset_id)

    def check_source_delete(self, source_id):
        if (
            self.db.one(
                "SELECT terminal_id FROM terminal_bindings WHERE source_id=?",
                (source_id,),
            )
            or self.db.one(
                "SELECT id FROM library_plans WHERE source_id=? AND archived=0",
                (source_id,),
            )
            or self.db.one(
                "SELECT id FROM source_images WHERE source_id=?", (source_id,)
            )
        ):
            raise FlowError(
                E.CONFLICT,
                "Remove terminal bindings, generation plans and local source pictures first",
                409,
            )