HalfikCode-ComfyUI-Prompt-DB/api_routes.py

"""
REST API routes for Prompt Database CRUD operations.
Registered with ComfyUI's PromptServer on import.
"""

import os
from aiohttp import web
from server import PromptServer
from . import db

HTML_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web")

routes = PromptServer.instance.routes


@routes.get("/prompt_db/manager")
async def serve_manager(request):
    """Serve the Prompt Database Manager HTML page."""
    html_path = os.path.join(HTML_DIR, "prompt_manager.html")
    return web.FileResponse(html_path)


@routes.get("/prompt_db/categories")
async def get_categories(request):
    categories = db.get_categories()
    return web.json_response({"categories": categories})


@routes.get("/prompt_db/prompts")
async def get_prompts(request):
    category = request.query.get("category", None)
    prompts = db.get_prompts(category)
    return web.json_response({"prompts": prompts})


@routes.get("/prompt_db/all")
async def get_all(request):
    data = db.load_db()
    return web.json_response(data)


@routes.get("/prompt_db/loras")
async def get_loras(request):
    """Return list of available LoRA files from ComfyUI's models folder."""
    loras = db.get_available_loras()
    return web.json_response({"loras": loras})


@routes.post("/prompt_db/prompt")
async def add_prompt(request):
    try:
        body = await request.json()
        category = body.get("category", "")
        name = body.get("name", "")
        text = body.get("text", "")
        loras = body.get("loras", [])

        if not category or not name:
            return web.json_response(
                {"error": "category and name are required"}, status=400
            )

        prompt = db.add_prompt(category, name, text, loras=loras)
        return web.json_response({"prompt": prompt})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.put("/prompt_db/prompt/{prompt_id}")
async def update_prompt(request):
    try:
        prompt_id = request.match_info["prompt_id"]
        body = await request.json()
        result = db.update_prompt(
            prompt_id,
            category=body.get("category"),
            name=body.get("name"),
            text=body.get("text"),
            loras=body.get("loras"),
        )
        if result:
            return web.json_response({"prompt": result})
        return web.json_response({"error": "Prompt not found"}, status=404)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.delete("/prompt_db/prompt/{prompt_id}")
async def delete_prompt(request):
    try:
        prompt_id = request.match_info["prompt_id"]
        if db.delete_prompt(prompt_id):
            return web.json_response({"success": True})
        return web.json_response({"error": "Prompt not found"}, status=404)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.post("/prompt_db/category")
async def add_category(request):
    try:
        body = await request.json()
        name = body.get("name", "")
        if not name:
            return web.json_response({"error": "name is required"}, status=400)
        if db.add_category(name):
            return web.json_response({"success": True})
        return web.json_response({"error": "Category already exists"}, status=409)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.delete("/prompt_db/category/{name}")
async def delete_category(request):
    try:
        name = request.match_info["name"]
        if db.delete_category(name):
            return web.json_response({"success": True})
        return web.json_response({"error": "Category not found"}, status=404)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.put("/prompt_db/categories/reorder")
async def reorder_categories(request):
    try:
        body = await request.json()
        ordered = body.get("categories", [])
        if not ordered:
            return web.json_response({"error": "categories list required"}, status=400)
        if db.reorder_categories(ordered):
            return web.json_response({"success": True})
        return web.json_response(
            {"error": "Category list does not match existing categories"}, status=400
        )
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.put("/prompt_db/prompts/reorder")
async def reorder_prompts(request):
    try:
        body = await request.json()
        ordered_ids = body.get("prompt_ids", [])
        if not ordered_ids:
            return web.json_response({"error": "prompt_ids list required"}, status=400)
        if db.reorder_prompts(ordered_ids):
            return web.json_response({"success": True})
        return web.json_response({"error": "Reorder failed"}, status=400)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.get("/prompt_db/export")
async def export_db_route(request):
    try:
        data = db.export_db()
        return web.json_response(data)
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)


@routes.post("/prompt_db/import")
async def import_db_route(request):
    try:
        body = await request.json()
        mode = body.get("mode", "merge")
        incoming = body.get("data")
        if not incoming or "categories" not in incoming or "prompts" not in incoming:
            return web.json_response(
                {"error": "Invalid import data: must have categories and prompts"},
                status=400,
            )
        stats = db.import_db(incoming, mode)
        return web.json_response({"success": True, **stats})
    except Exception as e:
        return web.json_response({"error": str(e)}, status=500)