"""
commands/workflow.py
====================

Workflow management commands for the Discord ComfyUI bot.

This module contains commands for loading and managing workflow templates.
"""

from __future__ import annotations

import json
import logging
from typing import Optional, Dict

from discord.ext import commands

from discord_utils import require_comfy_client


logger = logging.getLogger(__name__)


def setup_workflow_commands(bot):
    """
    Register workflow management commands with the bot.

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

    @bot.command(name="workflow-load", aliases=["workflowload", "wfl"], extras={"category": "Workflow"})
    @require_comfy_client
    async def load_workflow_command(ctx: commands.Context, *, path: Optional[str] = None) -> None:
        """
        Load a ComfyUI workflow from a JSON file.

        Usage:
            ttr!workflow-load path/to/workflow.json

        You can also attach a JSON file to the command message instead of
        providing a path.  The loaded workflow will replace the current
        workflow template used by the bot.  After loading a workflow you
        can generate images with your prompts while reusing the loaded
        graph structure.
        """
        workflow_data: Optional[Dict] = None

        # Check for attached JSON file first
        for attachment in ctx.message.attachments:
            if attachment.filename.lower().endswith(".json"):
                raw = await attachment.read()
                try:
                    text = raw.decode("utf-8")
                except UnicodeDecodeError as exc:
                    await ctx.reply(
                        f"`{attachment.filename}` is not valid UTF-8: {exc}",
                        mention_author=False,
                    )
                    return
                try:
                    workflow_data = json.loads(text)
                    break
                except json.JSONDecodeError as exc:
                    await ctx.reply(
                        f"Failed to parse `{attachment.filename}` as JSON: {exc}",
                        mention_author=False,
                    )
                    return

        # Otherwise try to load from provided path
        if workflow_data is None and path:
            try:
                with open(path, "r", encoding="utf-8") as f:
                    workflow_data = json.load(f)
            except FileNotFoundError:
                await ctx.reply(f"File not found: `{path}`", mention_author=False)
                return
            except json.JSONDecodeError as exc:
                await ctx.reply(f"Invalid JSON in `{path}`: {exc}", mention_author=False)
                return

        if workflow_data is None:
            await ctx.reply(
                "Please provide a JSON workflow file either as an attachment or a path.",
                mention_author=False,
            )
            return

        # Set the workflow on the client
        try:
            bot.comfy.set_workflow(workflow_data)
            await ctx.reply("Workflow loaded successfully.", mention_author=False)
        except Exception as exc:
            await ctx.reply(
                f"Failed to set workflow: {type(exc).__name__}: {exc}",
                mention_author=False,
            )