comfy-discord-web/commands/input_images.py

"""
commands/input_images.py
========================

Channel-backed input image management.

Images uploaded to the designated `comfy-input` channel get a persistent
"✅ Set as input" button posted by the bot — one reply per attachment so
every image in a multi-image message is independently selectable.

Persistent views survive bot restarts: on_ready re-registers every view
stored in the SQLite database.
"""

from __future__ import annotations

import io
import logging
from pathlib import Path

import discord
from discord.ext import commands

from image_utils import compress_to_discord_limit
from input_image_db import (
    activate_image_for_slot,
    get_all_images,
    upsert_image,
)

logger = logging.getLogger(__name__)

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp"}



class PersistentSetInputView(discord.ui.View):
    """
    A persistent view that survives bot restarts.

    One instance is created per DB row (i.e. per attachment).
    The button's custom_id encodes the row id so the callback can look
    up the exact filename to download.
    """

    def __init__(self, bot, config, row_id: int):
        super().__init__(timeout=None)
        self._bot = bot
        self._config = config
        self._row_id = row_id

        btn = discord.ui.Button(
            label="✅ Set as input",
            style=discord.ButtonStyle.success,
            custom_id=f"set_input:{row_id}",
        )
        btn.callback = self._set_callback
        self.add_item(btn)

    async def _set_callback(self, interaction: discord.Interaction) -> None:
        await interaction.response.defer(ephemeral=True)
        try:
            filename = activate_image_for_slot(
                self._row_id, "input_image", self._config.comfy_input_path
            )
            self._bot.comfy.state_manager.set_override("input_image", filename)
            await interaction.followup.send(
                f"✅ Input image set to `{filename}`", ephemeral=True
            )
        except Exception as exc:
            logger.exception("set_input button failed for row %s", self._row_id)
            await interaction.followup.send(f"❌ Error: {exc}", ephemeral=True)


async def _register_attachment(bot, config, message: discord.Message, attachment: discord.Attachment) -> None:
    """Post a reply with the image preview, a Set-as-input button, and record it in the DB."""
    logger.info("[_register_attachment] Start")
    original_data = await attachment.read()
    original_filename = attachment.filename
    logger.info("[_register_attachment] Reading attachment")

    # Compress only for the Discord re-send (8 MiB bot limit)
    send_data, send_filename = compress_to_discord_limit(original_data, original_filename)

    file = discord.File(io.BytesIO(send_data), filename=send_filename)
    reply = await message.channel.send(f"`{original_filename}`", file=file)

    # Store original quality bytes in DB
    row_id = upsert_image(message.id, reply.id, message.channel.id, original_filename, image_data=original_data)
    view = PersistentSetInputView(bot, config, row_id)
    bot.add_view(view, message_id=reply.id)
    logger.info("[_register_attachment] Done")
    await reply.edit(view=view)


def setup_input_image_commands(bot, config=None):
    """Register input image commands and the on_message listener."""

    @bot.listen("on_message")
    async def _on_input_channel_message(message: discord.Message) -> None:
        """Watch the comfy-input channel and attach a Set-as-input button to every image upload."""
        if config is None:
            logger.warning("[_on_input_channel_message] Config is none")
            return
        if message.channel.id != config.comfy_input_channel_id:
            return
        if message.author.bot:
            return

        image_attachments = [
            a for a in message.attachments
            if Path(a.filename).suffix.lower() in IMAGE_EXTENSIONS
        ]
        if not image_attachments:
            logger.info("[_on_input_channel_message] No image attachments")
            return

        for attachment in image_attachments:
            await _register_attachment(bot, config, message, attachment)

        try:
            await message.delete()
        except discord.Forbidden:
            logger.warning("Missing manage_messages permission to delete message %s", message.id)
        except Exception as exc:
            logger.warning("Could not delete message %s: %s", message.id, exc)

    @bot.command(
        name="sync-inputs",
        aliases=["si"],
        extras={"category": "Files"},
        help="Scan the comfy-input channel and add 'Set as input' buttons to any untracked images.",
    )
    async def sync_inputs_command(ctx: commands.Context) -> None:
        """Backfill Set-as-input buttons for images uploaded while the bot was offline."""
        if config is None:
            await ctx.reply("Bot config is not available.", mention_author=False)
            return

        channel = bot.get_channel(config.comfy_input_channel_id)
        if channel is None:
            try:
                channel = await bot.fetch_channel(config.comfy_input_channel_id)
            except Exception as exc:
                await ctx.reply(f"❌ Could not access input channel: {exc}", mention_author=False)
                return

        # Track existing records as (message_id, filename) pairs
        existing = {(row["original_message_id"], row["filename"]) for row in get_all_images()}

        new_count = 0
        async for message in channel.history(limit=None):
            if message.author.bot:
                continue

            had_new = False
            for attachment in message.attachments:
                if Path(attachment.filename).suffix.lower() not in IMAGE_EXTENSIONS:
                    continue
                if (message.id, attachment.filename) in existing:
                    continue

                await _register_attachment(bot, config, message, attachment)
                existing.add((message.id, attachment.filename))
                new_count += 1
                had_new = True

            if had_new:
                try:
                    await message.delete()
                except Exception as exc:
                    logger.warning("sync-inputs: could not delete message %s: %s", message.id, exc)

        already = len(get_all_images()) - new_count
        await ctx.reply(
            f"Synced {new_count} new image(s). {already} already known.",
            mention_author=False,
        )