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

import hashlib

from fastapi import Depends, File, Query, Request, UploadFile
from fastapi.responses import Response
from pydantic import Field

from pigway_frame_flow.core.auth import require
from pigway_frame_flow.core.photo_preferences import set_flags
from pigway_frame_flow.display.libraries import GenerationPlan
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.models import ImageExclusion, ImageStar, StrictModel
from pigway_frame_flow.providers.photo_sources import PhotoSourceInput
from pigway_frame_flow.providers.source_browser import BrowseFilters


class BindingInput(StrictModel):
    source_id: str = Field(min_length=1, max_length=120)
    plan_id: str = Field(pattern=r"^[a-f0-9]{32}$")


class UploadURLInput(StrictModel):
    url: str = Field(min_length=1, max_length=2048)


class ImportInput(StrictModel):
    item_id: str = Field(min_length=1, max_length=2048)


class RunInput(StrictModel):
    item_id: str | None = Field(default=None, min_length=1, max_length=2048)
    deliver: bool = False


class LibraryModeInput(StrictModel):
    enabled: bool


class GenerateInput(StrictModel):
    plan_id: str | None = Field(default=None, pattern=r"^[a-f0-9]{32}$")


def install_libraries(app, service):
    libraries, db = service.libraries, service.db

    @app.put(
        "/api/v1/jobs/{job_id}/library-mode",
        dependencies=[Depends(require("jobs:write"))],
    )
    def set_mode(job_id: str, payload: LibraryModeInput):
        return app.state.scheduler.set_library_mode(job_id, payload.enabled)

    @app.get("/api/v1/source-libraries", dependencies=[Depends(require("images:read"))])
    def sources():
        return [
            {
                **source,
                "image_count": db.one(
                    "SELECT COUNT(*) AS n FROM source_images WHERE source_id=?",
                    (source["id"],),
                )["n"],
            }
            for source in service.photo_sources.list()
        ]

    def save_upload(source_id, data, name):
        if len(data) > service.settings.max_bytes:
            raise FlowError(E.INVALID_REQUEST, "Image exceeds size limit")
        with service.lock:
            if source_id == "manual-uploads" and not db.one(
                "SELECT key FROM app_config WHERE key='photo_source:manual-uploads'"
            ):
                from pigway_frame_flow.renderers.images import convert

                convert(data, size=(1, 1))
                service.photo_sources.save(
                    source_id,
                    PhotoSourceInput(name="手动上传", kind="uploads", location="."),
                )
            if not service.photo_sources.get(source_id).enabled:
                raise FlowError(E.CONFLICT, "Photo source is disabled", 409)
            return libraries.ingest(
                source_id, "upload:" + hashlib.sha256(data).hexdigest(), data, name, {}
            )

    @app.post(
        "/api/v1/source-libraries/{source_id}/upload",
        dependencies=[Depends(require("images:write"))],
    )
    def upload(source_id: str, file: UploadFile = File(...)):
        return save_upload(
            source_id,
            file.file.read(service.settings.max_bytes + 1),
            file.filename or "upload.jpg",
        )

    @app.post(
        "/api/v1/source-libraries/{source_id}/upload-url",
        dependencies=[Depends(require("images:write"))],
    )
    def upload_url(source_id: str, payload: UploadURLInput):
        return save_upload(
            source_id, service.assets.fetch_url(payload.url), "url-upload.jpg"
        )

    @app.get(
        "/api/v1/source-libraries/{source_id}/images",
        dependencies=[Depends(require("images:read"))],
    )
    def pictures(
        source_id: str,
        offset: int = Query(0, ge=0),
        limit: int = Query(50, ge=1, le=200),
    ):
        service.photo_sources.get(source_id)
        rows = db.all(
            "SELECT id FROM source_images WHERE source_id=? ORDER BY created_at DESC,id DESC LIMIT ? OFFSET ?",
            (source_id, limit, offset),
        )
        return {
            "items": [libraries.picture(row["id"]) for row in rows],
            "offset": offset,
            "limit": limit,
            "total": db.one(
                "SELECT COUNT(*) AS n FROM source_images WHERE source_id=?",
                (source_id,),
            )["n"],
        }

    @app.get(
        "/api/v1/source-libraries/{source_id}/browse",
        dependencies=[Depends(require("images:read"))],
    )
    def browse(
        source_id: str,
        filters: BrowseFilters = Query(),
    ):
        return service.photo_sources.browser.page(
            source_id, filters.offset, filters.limit, filters
        )

    @app.get(
        "/api/v1/source-libraries/{source_id}/browse/{item_id}/preview",
        dependencies=[Depends(require("images:read"))],
    )
    def browse_preview(source_id: str, item_id: str):
        data = service.photo_sources.browser.preview(source_id, item_id)
        return Response(
            data,
            media_type="image/jpeg",
            headers={"Cache-Control": "private, no-store"},
        )

    @app.post(
        "/api/v1/source-libraries/{source_id}/import",
        dependencies=[Depends(require("admin"))],
    )
    def import_picture(source_id: str, payload: ImportInput):
        return libraries.import_picture(source_id, payload.item_id)

    @app.post(
        "/api/v1/source-libraries/{source_id}/run",
        dependencies=[Depends(require("admin"))],
    )
    def run(source_id: str, payload: RunInput, request: Request):
        with service.lock:
            return libraries.pipeline(
                source_id,
                payload.item_id,
                deliver=payload.deliver,
                caller=request.state.client_id,
            )

    @app.get("/api/v1/generation-plans", dependencies=[Depends(require("images:read"))])
    def plans(source_id: str):
        return libraries.plans(source_id)

    @app.post("/api/v1/generation-plans", dependencies=[Depends(require("admin"))])
    def create_plan(payload: GenerationPlan):
        return libraries.save_plan(payload)

    @app.put(
        "/api/v1/generation-plans/{plan_id}", dependencies=[Depends(require("admin"))]
    )
    def update_plan(plan_id: str, payload: GenerationPlan):
        return libraries.save_plan(payload, plan_id)

    @app.delete(
        "/api/v1/generation-plans/{plan_id}",
        status_code=204,
        dependencies=[Depends(require("admin"))],
    )
    def delete_plan(plan_id: str):
        libraries.archive(plan_id)

    @app.post(
        "/api/v1/source-images/{picture_id}/generate",
        dependencies=[Depends(require("admin"))],
    )
    def generate(picture_id: str, payload: GenerateInput):
        with service.lock:
            picture = libraries.picture(picture_id)
            ids = (
                [payload.plan_id]
                if payload.plan_id
                else [
                    p["id"]
                    for p in libraries.plans(picture["source_id"])
                    if not p["archived"]
                ]
            )
            return [libraries.generate(picture_id, plan_id) for plan_id in ids]

    @app.post(
        "/api/v1/source-images/{picture_id}/preview",
        dependencies=[Depends(require("images:read"))],
    )
    def preview(picture_id: str, payload: GenerationPlan):
        picture = libraries.picture(picture_id)
        if picture["source_id"] != payload.source_id:
            raise FlowError(
                E.CONFLICT, "Picture and plan must share a source library", 409
            )
        data = libraries.render(
            service.assets.path(picture_id).read_bytes(), picture["caption"], payload
        )
        return Response(
            data, media_type="image/jpeg", headers={"Cache-Control": "no-store"}
        )

    @app.put(
        "/api/v1/source-images/{picture_id}/star",
        dependencies=[Depends(require("images:write"))],
    )
    def star(picture_id: str, payload: ImageStar):
        picture = libraries.picture(picture_id)
        return set_flags(
            db, picture["source_id"], picture["item_id"], starred=payload.starred
        )

    @app.put(
        "/api/v1/source-images/{picture_id}/exclusion",
        dependencies=[Depends(require("images:write"))],
    )
    def exclude(picture_id: str, payload: ImageExclusion):
        picture = libraries.picture(picture_id)
        return set_flags(
            db, picture["source_id"], picture["item_id"], excluded=payload.excluded
        )

    @app.delete(
        "/api/v1/source-images/{picture_id}",
        status_code=204,
        dependencies=[Depends(require("images:delete"))],
    )
    def delete_picture(picture_id: str):
        libraries.delete_picture(picture_id)

    @app.delete(
        "/api/v1/library-versions/{asset_id}",
        status_code=204,
        dependencies=[Depends(require("images:delete"))],
    )
    def delete_version(asset_id: str):
        with service.lock:
            if not db.one(
                "SELECT asset_id FROM library_versions WHERE asset_id=?", (asset_id,)
            ):
                raise FlowError(E.NOT_FOUND, "Generated image not found", 404)
            libraries.check_asset_references(asset_id)
            libraries.remove_asset(asset_id)

    @app.get(
        "/api/v1/library-terminals", dependencies=[Depends(require("images:read"))]
    )
    def terminals():
        return libraries.terminals()

    @app.put(
        "/api/v1/library-terminals/{terminal_id}/binding",
        dependencies=[Depends(require("admin"))],
    )
    def bind(terminal_id: str, payload: BindingInput):
        return libraries.bind(terminal_id, payload.source_id, payload.plan_id)

    @app.delete(
        "/api/v1/library-terminals/{terminal_id}/binding",
        status_code=204,
        dependencies=[Depends(require("admin"))],
    )
    def unbind(terminal_id: str):
        with service.lock:
            db.execute(
                "DELETE FROM terminal_bindings WHERE terminal_id=?", (terminal_id,)
            )

    @app.get(
        "/api/v1/library-terminals/{terminal_id}/images",
        dependencies=[Depends(require("images:read"))],
    )
    def terminal_images(
        terminal_id: str,
        offset: int = Query(0, ge=0),
        limit: int = Query(50, ge=1, le=200),
    ):
        terminal = next(
            (t for t in libraries.terminals() if t["id"] == terminal_id), None
        )
        if terminal is None:
            raise FlowError(E.NOT_FOUND, "Terminal not found", 404)
        binding = libraries.binding(terminal_id)
        if not binding:
            return {"items": [], "total": 0, "offset": offset, "limit": limit}
        where = "FROM library_versions v JOIN source_images s ON s.id=v.source_image_id WHERE v.plan_id=? AND s.source_id=?"
        args = (binding["plan_id"], binding["source_id"])
        rows = db.all(
            "SELECT v.asset_id,v.source_image_id "
            + where
            + " ORDER BY v.created_at DESC,v.asset_id LIMIT ? OFFSET ?",
            (*args, limit, offset),
        )
        current_asset = None
        if rows and terminal["transport_supported"]:
            try:
                device = service.registry.get(terminal_id).device()
                if device.status == "online":
                    current = service.current_content(terminal_id, device=device)
                    if current["synced"] and current["source_known"]:
                        current_asset = current["asset_id"]
            except FlowError:
                # Offline devices must not hide locally generated versions.
                pass
        items = []
        versions_by_picture = {}
        for row in rows:
            picture_id = row["source_image_id"]
            if picture_id not in versions_by_picture:
                versions_by_picture[picture_id] = {
                    v["asset_id"]: v for v in libraries.versions(picture_id)
                }
            version = versions_by_picture[picture_id][row["asset_id"]]
            delivery = db.one(
                "SELECT operation_id FROM terminal_deliveries WHERE terminal_id=? AND asset_id=? ORDER BY created_at DESC,operation_id DESC LIMIT 1",
                (terminal_id, row["asset_id"]),
            )
            operation = (
                db.one(
                    "SELECT id AS operation_id,status,created_at,updated_at FROM operations WHERE id=?",
                    (delivery["operation_id"],),
                )
                if delivery and delivery["operation_id"]
                else None
            )
            items.append(
                {
                    **version,
                    "operation": operation,
                    "current": row["asset_id"] == current_asset,
                    "delivery_status": operation["status"]
                    if operation
                    else "GENERATED",
                }
            )
        return {
            "items": items,
            "total": db.one("SELECT COUNT(*) AS n " + where, args)["n"],
            "offset": offset,
            "limit": limit,
        }

    @app.post(
        "/api/v1/library-terminals/{terminal_id}/images/{asset_id}/send",
        dependencies=[Depends(require("display:write"))],
    )
    def send(terminal_id: str, asset_id: str, request: Request):
        return {
            "operation_id": libraries.send(
                terminal_id, asset_id, client_id=request.state.client_id
            )
        }