comfy-discord-web/commands/history.py

"""
commands/history.py
===================

History management commands for the Discord ComfyUI bot.

This module contains commands for viewing and retrieving past generation
results from the bot's history.
"""

from __future__ import annotations

import logging
from io import BytesIO
from typing import Optional

import discord
from discord.ext import commands

from config import MAX_IMAGES_PER_RESPONSE
from discord_utils import require_comfy_client, truncate_text, convert_image_bytes_to_discord_files


logger = logging.getLogger(__name__)


def setup_history_commands(bot):
    """
    Register history management commands with the bot.

    Parameters
    ----------
    bot : commands.Bot
        The Discord bot instance.
    """

    @bot.command(name="history", extras={"category": "History"})
    @require_comfy_client
    async def history_command(ctx: commands.Context) -> None:
        """
        Show a list of recently generated prompts.

        The bot keeps a rolling history of the last few generations.  Each
        entry lists the prompt id along with the positive and negative
        prompt texts.  You can retrieve the images from a previous
        generation with the ``ttr!gethistory <prompt_id>`` command.
        """
        hist = bot.comfy.get_history()
        if not hist:
            await ctx.reply(
                "No history available yet.  Generate something first!",
                mention_author=False,
            )
            return

        # Build a human readable list
        lines = ["Here are the most recent generations (oldest first):"]
        for entry in hist:
            pid = entry.get("prompt_id", "unknown")
            prompt = entry.get("prompt") or ""
            neg = entry.get("negative_prompt") or ""
            # Truncate long prompts for readability
            lines.append(
                f"• ID: {pid} | prompt: '{truncate_text(prompt, 60)}' | negative: '{truncate_text(neg, 60)}'"
            )
        await ctx.reply("\n".join(lines), mention_author=False)

    @bot.command(name="get-history", aliases=["gethistory", "gh"], extras={"category": "History"})
    @require_comfy_client
    async def get_history_command(ctx: commands.Context, *, arg: str = "") -> None:
        """
        Retrieve images from a previous generation, or search history by keyword.

        Usage:
            ttr!gethistory <prompt_id_or_index>
            ttr!gethistory search:<keyword>

        Provide either the prompt id returned in the generation response
        (shown in `ttr!history`) or the 1‑based index into the history
        list.  The bot will fetch the images associated with that
        generation and resend them.  If no images are found, you will be
        notified.

        Use ``search:<keyword>`` to filter history by prompt text, checkpoint
        name, seed value, or any other override field.
        """
        if not arg:
            await ctx.reply(
                "Please provide a prompt id, history index, or `search:<keyword>`.  See `ttr!history` for a list.",
                mention_author=False,
            )
            return

        # Handle search:<keyword>
        lower_arg = arg.lower()
        if lower_arg.startswith("search:"):
            keyword = arg[len("search:"):].strip()
            if not keyword:
                await ctx.reply("Please provide a keyword after `search:`.", mention_author=False)
                return
            from generation_db import search_history_for_user, get_history as db_get_history
            # Use get_history for Discord since Discord bot doesn't have per-user context like the web UI
            hist = db_get_history(limit=50)
            matches = [
                e for e in hist
                if keyword.lower() in str(e.get("overrides", {})).lower()
            ]
            if not matches:
                await ctx.reply(f"No history entries matching `{keyword}`.", mention_author=False)
                return
            lines = [f"**History matching `{keyword}`** ({len(matches)} result(s))"]
            for entry in matches[:10]:
                pid = entry.get("prompt_id", "unknown")
                overrides = entry.get("overrides") or {}
                prompt = str(overrides.get("prompt") or "")
                lines.append(
                    f"• `{pid[:12]}…` | {truncate_text(prompt, 60) if prompt else '(no prompt)'}"
                )
            if len(matches) > 10:
                lines.append(f"_(showing first 10 of {len(matches)})_")
            await ctx.reply("\n".join(lines), mention_author=False)
            return

        # Determine whether arg refers to an index or an id
        target_id: Optional[str] = None
        hist = bot.comfy.get_history()

        # If arg is a digit, interpret as 1‑based index
        if arg.isdigit():
            idx = int(arg) - 1
            if idx < 0 or idx >= len(hist):
                await ctx.reply(
                    f"Index out of range.  There are {len(hist)} entries in history.",
                    mention_author=False,
                )
                return
            target_id = hist[idx]["prompt_id"]
        else:
            # Otherwise treat as an explicit prompt id
            target_id = arg.strip()

        try:
            images = await bot.comfy.fetch_history_images(target_id)
            if not images:
                await ctx.reply(
                    f"No images found for prompt id `{target_id}`.",
                    mention_author=False,
                )
                return

            files = []
            for idx, img_bytes in enumerate(images):
                if idx >= MAX_IMAGES_PER_RESPONSE:
                    break
                file_obj = BytesIO(img_bytes)
                file_obj.seek(0)
                files.append(discord.File(file_obj, filename=f"history_{target_id}_{idx+1}.png"))

            await ctx.reply(
                content=f"Here are the images for prompt id `{target_id}`:",
                files=files,
                mention_author=False,
            )
        except Exception as exc:
            logger.exception("Failed to fetch history for %s", target_id)
            await ctx.reply(
                f"An error occurred: {type(exc).__name__}: {exc}",
                mention_author=False,
            )