comfy-discord-web/web/routers/inputs_router.py

"""GET/POST/DELETE /api/inputs; GET /api/inputs/{id}/image; POST /api/inputs/{id}/activate"""
from __future__ import annotations

import logging
import mimetypes
from pathlib import Path
from typing import Optional

from fastapi import APIRouter, Body, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import Response

from web.auth import require_auth
from web.deps import get_config, get_user_registry

router = APIRouter()
logger = logging.getLogger(__name__)


@router.get("")
async def list_inputs(_: dict = Depends(require_auth)):
    """List all input images (Discord + web uploads)."""
    from input_image_db import get_all_images
    rows = get_all_images()
    return [dict(r) for r in rows]


@router.post("")
async def upload_input(
    file: UploadFile = File(...),
    slot_key: Optional[str] = Form(default=None),
    user: dict = Depends(require_auth),
):
    """
    Upload an input image.

    Stores image bytes directly in SQLite.  If *slot_key* is provided the
    image is immediately activated for that slot (writes to ComfyUI input
    folder and updates the user's state override).

    The physical slot file uses a namespaced key ``<user_label>_<slot_key>``
    so concurrent users each get their own active image file.
    """
    config = get_config()
    if config is None:
        raise HTTPException(503, "Config not available")

    data = await file.read()
    filename = file.filename or "upload.png"

    from input_image_db import upsert_image, activate_image_for_slot
    row_id = upsert_image(
        original_message_id=0,  # sentinel for web uploads
        bot_reply_id=0,
        channel_id=0,
        filename=filename,
        image_data=data,
    )

    activated_filename: str | None = None
    if slot_key:
        user_label: str = user["sub"]
        namespaced_key = f"{user_label}_{slot_key}"
        activated_filename = activate_image_for_slot(
            row_id, namespaced_key, config.comfy_input_path
        )
        registry = get_user_registry()
        if registry:
            registry.get_state_manager(user_label).set_override(slot_key, activated_filename)
        else:
            from web.deps import get_comfy
            comfy = get_comfy()
            if comfy:
                comfy.state_manager.set_override(slot_key, activated_filename)

    return {"id": row_id, "filename": filename, "slot_key": slot_key, "activated_filename": activated_filename}


@router.post("/{row_id}/activate")
async def activate_input(
    row_id: int,
    slot_key: str = Body(default="input_image", embed=True),
    user: dict = Depends(require_auth),
):
    """Write the stored image to the ComfyUI input folder and set the user's slot override."""
    config = get_config()
    if config is None:
        raise HTTPException(503, "Config not available")

    from input_image_db import get_image, activate_image_for_slot
    row = get_image(row_id)
    if row is None:
        raise HTTPException(404, "Image not found")

    user_label: str = user["sub"]
    namespaced_key = f"{user_label}_{slot_key}"

    try:
        filename = activate_image_for_slot(row_id, namespaced_key, config.comfy_input_path)
    except ValueError as exc:
        raise HTTPException(409, str(exc))

    registry = get_user_registry()
    if registry:
        registry.get_state_manager(user_label).set_override(slot_key, filename)
    else:
        from web.deps import get_comfy
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "State manager not available")
        comfy.state_manager.set_override(slot_key, filename)

    return {"ok": True, "slot_key": slot_key, "filename": filename}


@router.delete("/{row_id}")
async def delete_input(row_id: int, _: dict = Depends(require_auth)):
    """Delete an input image record (and its active slot file if applicable)."""
    from input_image_db import get_image, delete_image
    row = get_image(row_id)
    if row is None:
        raise HTTPException(404, "Image not found")

    config = get_config()
    delete_image(row_id, comfy_input_path=config.comfy_input_path if config else None)
    return {"ok": True}


@router.get("/{row_id}/image")
async def get_input_image(row_id: int, _: dict = Depends(require_auth)):
    """Serve the raw image bytes stored in the database for a given input image row."""
    from input_image_db import get_image, get_image_data
    row = get_image(row_id)
    if row is None:
        raise HTTPException(404, "Image not found")

    data = get_image_data(row_id)
    if data is None:
        raise HTTPException(404, "Image data not available — re-upload to backfill")

    mime, _ = mimetypes.guess_type(row["filename"])
    return Response(content=data, media_type=mime or "application/octet-stream")


def _pil_resize_response(data: bytes, filename: str, max_size: int, quality: int) -> Response:
    """Resize image bytes with Pillow and return a JPEG Response. Raises on failure."""
    import io
    from PIL import Image as _PIL
    img = _PIL.open(io.BytesIO(data))
    img.thumbnail((max_size, max_size), _PIL.LANCZOS)
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=quality, optimize=True)
    return Response(
        content=buf.getvalue(),
        media_type="image/jpeg",
        headers={"Cache-Control": "public, max-age=86400"},
    )


@router.get("/{row_id}/thumb")
async def get_input_thumb(row_id: int, _: dict = Depends(require_auth)):
    """Serve a small compressed thumbnail (max 200 px, JPEG 65 %) for fast previews."""
    from input_image_db import get_image, get_image_data
    row = get_image(row_id)
    if row is None:
        raise HTTPException(404, "Image not found")

    data = get_image_data(row_id)
    if data is None:
        raise HTTPException(404, "Image data not available — re-upload to backfill")

    try:
        return _pil_resize_response(data, row["filename"], max_size=200, quality=65)
    except Exception:
        mime, _ = mimetypes.guess_type(row["filename"])
        return Response(content=data, media_type=mime or "application/octet-stream")


@router.get("/{row_id}/mid")
async def get_input_mid(row_id: int, _: dict = Depends(require_auth)):
    """Serve a medium compressed image (max 800 px, JPEG 80 %) for progressive loading."""
    from input_image_db import get_image, get_image_data
    row = get_image(row_id)
    if row is None:
        raise HTTPException(404, "Image not found")

    data = get_image_data(row_id)
    if data is None:
        raise HTTPException(404, "Image data not available — re-upload to backfill")

    try:
        return _pil_resize_response(data, row["filename"], max_size=800, quality=80)
    except Exception:
        mime, _ = mimetypes.guess_type(row["filename"])
        return Response(content=data, media_type=mime or "application/octet-stream")