# PromptManager/py/api.py

import datetime
import json
import os
import traceback
from typing import Any, Dict, List, Optional

import server
from aiohttp import web
from PIL import Image

# Import database operations
try:
    from ..database.operations import PromptDatabase
    from ..utils.logging_config import get_logger
except ImportError:
    # Fallback for when module isn't imported as package
    import os
    import sys

    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from database.operations import PromptDatabase
    from utils.logging_config import get_logger


class PromptManagerAPI:
    """API class for PromptManager database operations."""

    def __init__(self):
        self.logger = get_logger('prompt_manager.api')
        self.logger.info("Initializing PromptManager API")
        
        self.db = PromptDatabase()

        # Run cleanup on initialization to remove any existing duplicates
        try:
            removed = self.db.cleanup_duplicates()
            if removed > 0:
                self.logger.info(f"Startup cleanup: removed {removed} duplicate prompts")
        except Exception as e:
            self.logger.error(f"Startup cleanup failed: {e}")
        
        self.logger.info("PromptManager API initialization completed")

    def add_routes(self, routes):
        """Add API routes to ComfyUI server using decorator pattern."""

        # Test route to verify registration works
        @routes.get("/prompt_manager/test")
        async def test_route(request):
            return web.json_response(
                {
                    "success": True,
                    "message": "PromptManager API is working!",
                    "timestamp": str(datetime.datetime.now()),
                }
            )

        @routes.get("/prompt_manager/search")
        async def search_prompts_route(request):
            return await self.search_prompts(request)

        @routes.get("/prompt_manager/recent")
        async def get_recent_prompts_route(request):
            return await self.get_recent_prompts(request)

        @routes.get("/prompt_manager/categories")
        async def get_categories_route(request):
            return await self.get_categories(request)

        @routes.get("/prompt_manager/tags")
        async def get_tags_route(request):
            return await self.get_tags(request)

        @routes.post("/prompt_manager/cleanup")
        async def cleanup_duplicates_route(request):
            return await self.cleanup_duplicates_endpoint(request)

        @routes.post("/prompt_manager/maintenance")
        async def maintenance_route(request):
            return await self.run_maintenance(request)

        @routes.post("/prompt_manager/save")
        async def save_prompt_route(request):
            return await self.save_prompt(request)

        @routes.delete("/prompt_manager/delete/{prompt_id}")
        async def delete_prompt_route(request):
            return await self.delete_prompt(request)

        # Serve the web UI HTML file
        @routes.get("/prompt_manager/web")
        async def serve_web_ui(request):
            try:
                import os

                # Get the path to the web directory
                current_dir = os.path.dirname(
                    os.path.dirname(os.path.abspath(__file__))
                )
                html_path = os.path.join(current_dir, "web", "index.html")

                if os.path.exists(html_path):
                    with open(html_path, "r", encoding="utf-8") as f:
                        html_content = f.read()

                    return web.Response(
                        text=html_content, content_type="text/html", charset="utf-8"
                    )
                else:
                    return web.Response(
                        text="<h1>Web UI not found</h1><p>HTML file not located at expected path.</p>",
                        content_type="text/html",
                        status=404,
                    )

            except Exception as e:
                return web.Response(
                    text=f"<h1>Error</h1><p>Failed to load web UI: {str(e)}</p>",
                    content_type="text/html",
                    status=500,
                )

        # Serve the gallery interface
        @routes.get("/prompt_manager/gallery.html")
        async def serve_gallery_ui(request):
            try:
                import os

                current_dir = os.path.dirname(
                    os.path.dirname(os.path.abspath(__file__))
                )
                html_path = os.path.join(current_dir, "web", "metadata.html")

                if os.path.exists(html_path):
                    with open(html_path, "r", encoding="utf-8") as f:
                        html_content = f.read()

                    return web.Response(
                        text=html_content, content_type="text/html", charset="utf-8"
                    )
                else:
                    return web.Response(
                        text="<h1>Gallery not found</h1><p>gallery.html file not located at expected path.</p>",
                        content_type="text/html",
                        status=404,
                    )

            except Exception as e:
                return web.Response(
                    text=f"<h1>Error</h1><p>Failed to load gallery: {str(e)}</p>",
                    content_type="text/html",
                    status=500,
                )

        # Serve the admin interface
        @routes.get("/prompt_manager/admin")
        async def serve_admin_ui(request):
            try:
                import os

                current_dir = os.path.dirname(
                    os.path.dirname(os.path.abspath(__file__))
                )
                html_path = os.path.join(current_dir, "web", "admin.html")

                if os.path.exists(html_path):
                    with open(html_path, "r", encoding="utf-8") as f:
                        html_content = f.read()

                    return web.Response(
                        text=html_content, content_type="text/html", charset="utf-8"
                    )
                else:
                    return web.Response(
                        text="<h1>Admin UI not found</h1>",
                        content_type="text/html",
                        status=404,
                    )

            except Exception as e:
                return web.Response(
                    text=f"<h1>Error</h1><p>Failed to load admin UI: {str(e)}</p>",
                    content_type="text/html",
                    status=500,
                )

        # Statistics endpoint
        @routes.get("/prompt_manager/stats")
        async def get_stats_route(request):
            return await self.get_statistics(request)

        # Settings endpoints
        @routes.get("/prompt_manager/settings")
        async def get_settings_route(request):
            return await self.get_settings(request)

        @routes.post("/prompt_manager/settings")
        async def save_settings_route(request):
            return await self.save_settings(request)


        # Individual prompt management
        @routes.put("/prompt_manager/prompts/{prompt_id}")
        async def update_prompt_route(request):
            return await self.update_prompt(request)

        @routes.put("/prompt_manager/prompts/{prompt_id}/rating")
        async def update_rating_route(request):
            return await self.update_prompt_rating(request)

        @routes.post("/prompt_manager/prompts/{prompt_id}/tags")
        async def add_tag_route(request):
            return await self.add_prompt_tag(request)

        @routes.delete("/prompt_manager/prompts/{prompt_id}/tags")
        async def remove_tag_route(request):
            return await self.remove_prompt_tag(request)

        @routes.post("/prompt_manager/prompts/tags")
        async def add_tags_to_prompt_route(request):
            return await self.add_tags_to_prompt(request)

        # Bulk operations
        @routes.post("/prompt_manager/bulk/delete")
        async def bulk_delete_route(request):
            return await self.bulk_delete_prompts(request)

        @routes.post("/prompt_manager/bulk/tags")
        async def bulk_add_tags_route(request):
            return await self.bulk_add_tags(request)

        @routes.post("/prompt_manager/bulk/category")
        async def bulk_set_category_route(request):
            return await self.bulk_set_category(request)

        # Export functionality
        @routes.get("/prompt_manager/export")
        async def export_prompts_route(request):
            return await self.export_prompts(request)

        # Database backup and restore
        @routes.get("/prompt_manager/backup")
        async def backup_database_route(request):
            return await self.backup_database(request)

        @routes.post("/prompt_manager/restore")
        async def restore_database_route(request):
            return await self.restore_database(request)

        # Gallery endpoints
        @routes.get("/prompt_manager/prompts/{prompt_id}/images")
        async def get_prompt_images_route(request):
            return await self.get_prompt_images(request)

        @routes.get("/prompt_manager/images/recent")
        async def get_recent_images_route(request):
            return await self.get_recent_images(request)

        @routes.get("/prompt_manager/images/search")
        async def search_images_route(request):
            return await self.search_images(request)

        @routes.get("/prompt_manager/images/{image_id}/file")
        async def serve_image_route(request):
            return await self.serve_image(request)

        @routes.post("/prompt_manager/images/link")
        async def link_image_route(request):
            return await self.link_image_to_prompt(request)

        @routes.delete("/prompt_manager/images/{image_id}")
        async def delete_image_route(request):
            return await self.delete_image(request)

        # Diagnostic endpoints
        @routes.get("/prompt_manager/diagnostics")
        async def run_diagnostics_route(request):
            return await self.run_diagnostics(request)

        @routes.post("/prompt_manager/diagnostics/test-link")
        async def test_image_link_route(request):
            return await self.test_image_link(request)

        # Scan endpoint
        @routes.post("/prompt_manager/scan")
        async def scan_images_route(request):
            return await self.scan_images(request)

        # Logging endpoints
        @routes.get("/prompt_manager/logs")
        async def get_logs_route(request):
            return await self.get_logs(request)

        @routes.get("/prompt_manager/logs/files")
        async def get_log_files_route(request):
            return await self.get_log_files(request)

        @routes.get("/prompt_manager/logs/download/{filename}")
        async def download_log_route(request):
            return await self.download_log_file(request)

        @routes.post("/prompt_manager/logs/truncate")
        async def truncate_logs_route(request):
            return await self.truncate_logs(request)

        @routes.get("/prompt_manager/logs/config")
        async def get_log_config_route(request):
            return await self.get_log_config(request)

        @routes.post("/prompt_manager/logs/config")
        async def update_log_config_route(request):
            return await self.update_log_config(request)

        @routes.get("/prompt_manager/logs/stats")
        async def get_log_stats_route(request):
            return await self.get_log_stats(request)

        self.logger.info("All routes registered with decorator pattern")

    async def search_prompts(self, request):
        """
        Search prompts endpoint.
        GET /prompt_manager/search?text=...&category=...&tags=...&min_rating=...&limit=...
        """
        try:
            # Get query parameters
            text = request.query.get("text", "").strip()
            category = request.query.get("category", "").strip()
            tags_str = request.query.get("tags", "").strip()
            min_rating = request.query.get("min_rating", 0)
            limit = int(request.query.get("limit", 50))

            # Parse tags
            tags = None
            if tags_str:
                tags = [tag.strip() for tag in tags_str.split(",") if tag.strip()]

            # Parse min_rating
            try:
                min_rating = int(min_rating) if min_rating else None
            except ValueError:
                min_rating = None

            # Perform search
            results = self.db.search_prompts(
                text=text if text else None,
                category=category if category else None,
                tags=tags,
                rating_min=min_rating,
                limit=limit,
            )

            return web.json_response(
                {"success": True, "results": results, "count": len(results)}
            )

        except Exception as e:
            self.logger.error(f"Search error: {e}", exc_info=True)
            return web.json_response(
                {"success": False, "error": f"Search failed: {str(e)}", "results": []},
                status=500,
            )

    async def get_recent_prompts(self, request):
        """
        Get recent prompts endpoint with pagination support.
        GET /prompt_manager/recent?limit=50&page=2&offset=100
        """
        try:
            limit = int(request.query.get("limit", 50))
            page = int(request.query.get("page", 1))
            offset = int(request.query.get("offset", 0))
            
            # If page is provided, calculate offset from page
            if page > 1 and offset == 0:
                offset = (page - 1) * limit
            
            # Ensure reasonable limits
            if limit > 1000:
                limit = 1000
            elif limit < 1:
                limit = 1

            results = self.db.get_recent_prompts(limit=limit, offset=offset)

            return web.json_response({
                "success": True, 
                "results": results['prompts'],
                "pagination": {
                    "total": results['total'],
                    "limit": results['limit'],
                    "offset": results['offset'],
                    "page": results['page'],
                    "total_pages": results['total_pages'],
                    "has_more": results['has_more'],
                    "count": len(results['prompts'])
                }
            })

        except Exception as e:
            self.logger.error(f"Recent prompts error: {e}", exc_info=True)
            return web.json_response(
                {
                    "success": False,
                    "error": f"Failed to get recent prompts: {str(e)}",
                    "results": [],
                    "pagination": {"total": 0, "page": 1, "total_pages": 0}
                },
                status=500,
            )

    async def get_categories(self, request):
        """
        Get all categories endpoint.
        GET /prompt_manager/categories
        """
        try:
            categories = self.db.get_all_categories()

            return web.json_response({"success": True, "categories": categories})

        except Exception as e:
            self.logger.error(f"Categories error: {e}")
            return web.json_response(
                {
                    "success": False,
                    "error": f"Failed to get categories: {str(e)}",
                    "categories": [],
                },
                status=500,
            )

    async def get_tags(self, request):
        """
        Get all tags endpoint.
        GET /prompt_manager/tags
        """
        try:
            tags = self.db.get_all_tags()

            return web.json_response({"success": True, "tags": tags})

        except Exception as e:
            self.logger.error(f"Tags error: {e}")
            return web.json_response(
                {
                    "success": False,
                    "error": f"Failed to get tags: {str(e)}",
                    "tags": [],
                },
                status=500,
            )

    async def save_prompt(self, request):
        """
        Save a new prompt endpoint.
        POST /prompt_manager/save
        Body: {"text": "...", "category": "...", "tags": [...], "rating": 5, "notes": "..."}
        """
        try:
            data = await request.json()

            text = data.get("text", "").strip()
            if not text:
                return web.json_response(
                    {"success": False, "error": "Text is required"}, status=400
                )

            category = data.get("category", "").strip() or None
            tags = data.get("tags", [])
            rating = data.get("rating") or None
            notes = data.get("notes", "").strip() or None

            # Generate hash for duplicate detection
            import hashlib

            prompt_hash = hashlib.sha256(text.encode("utf-8")).hexdigest()
            
            # Check if prompt already exists
            existing = self.db.get_prompt_by_hash(prompt_hash)
            if existing:
                # Update metadata if this is a duplicate with new info
                if any([category, tags, rating, notes]):
                    self.db.update_prompt_metadata(
                        prompt_id=existing['id'],
                        category=category,
                        tags=tags,
                        rating=rating,
                        notes=notes
                    )
                return web.json_response(
                    {
                        "success": True,
                        "prompt_id": existing['id'],
                        "message": "Prompt already exists, metadata updated",
                        "is_duplicate": True
                    }
                )

            # Save new prompt
            prompt_id = self.db.save_prompt(
                text=text,
                category=category,
                tags=tags if tags else None,
                rating=rating,
                notes=notes,
                prompt_hash=prompt_hash,
            )

            return web.json_response(
                {
                    "success": True,
                    "prompt_id": prompt_id,
                    "message": "Prompt saved successfully",
                }
            )

        except Exception as e:
            self.logger.error(f"Save error: {e}", exc_info=True)
            return web.json_response(
                {"success": False, "error": f"Failed to save prompt: {str(e)}"},
                status=500,
            )

    async def delete_prompt(self, request):
        """
        Delete a prompt endpoint.
        DELETE /prompt_manager/delete/{prompt_id}
        """
        try:
            prompt_id = int(request.match_info["prompt_id"])

            success = self.db.delete_prompt(prompt_id)

            if success:
                return web.json_response(
                    {"success": True, "message": "Prompt deleted successfully"}
                )
            else:
                return web.json_response(
                    {
                        "success": False,
                        "error": "Prompt not found or could not be deleted",
                    },
                    status=404,
                )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Delete error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to delete prompt: {str(e)}"},
                status=500,
            )

    async def cleanup_duplicates_endpoint(self, request):
        """
        Cleanup duplicate prompts endpoint.
        POST /prompt_manager/cleanup
        """
        try:
            removed_count = self.db.cleanup_duplicates()

            return web.json_response(
                {
                    "success": True,
                    "message": f"Cleanup completed",
                    "duplicates_removed": removed_count,
                }
            )

        except Exception as e:
            self.logger.error(f"Cleanup error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to cleanup duplicates: {str(e)}"},
                status=500,
            )

    async def get_statistics(self, request):
        """Get database statistics."""
        try:
            # Get basic stats from database
            with self.db.model.get_connection() as conn:
                cursor = conn.execute("SELECT COUNT(*) as total FROM prompts")
                total_prompts = cursor.fetchone()["total"]

                cursor = conn.execute(
                    "SELECT COUNT(DISTINCT TRIM(category)) as total FROM prompts WHERE category IS NOT NULL AND TRIM(category) != ''"
                )
                total_categories = cursor.fetchone()["total"]

                cursor = conn.execute(
                    "SELECT AVG(rating) as avg FROM prompts WHERE rating IS NOT NULL"
                )
                avg_rating = cursor.fetchone()["avg"]

                # Count unique tags
                cursor = conn.execute("SELECT tags FROM prompts WHERE tags IS NOT NULL")
                all_tags = set()
                for row in cursor.fetchall():
                    try:
                        tags = json.loads(row["tags"])
                        if isinstance(tags, list):
                            all_tags.update(tags)
                    except:
                        continue

                # Get all categories for debugging (including raw data)
                cursor = conn.execute(
                    "SELECT DISTINCT category FROM prompts WHERE category IS NOT NULL ORDER BY category"
                )
                raw_categories = [row['category'] for row in cursor.fetchall()]
                
                cursor = conn.execute(
                    "SELECT DISTINCT TRIM(category) as category FROM prompts WHERE category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
                )
                filtered_categories = [row['category'] for row in cursor.fetchall()]
                
                # Debug logging with detailed category info
                self.logger.debug(f"Statistics calculated - Prompts: {total_prompts}, Categories: {total_categories}, Tags: {len(all_tags)}, Avg Rating: {avg_rating}")
                self.logger.debug(f"Raw categories from DB: {[(cat, len(cat), repr(cat)) for cat in raw_categories]}")
                self.logger.debug(f"Filtered categories: {[(cat, len(cat), repr(cat)) for cat in filtered_categories]}")

                return web.json_response(
                    {
                        "success": True,
                        "stats": {
                            "total_prompts": total_prompts,
                            "total_categories": total_categories,
                            "unique_categories": total_categories,  # Keep both for compatibility
                            "total_tags": len(all_tags),
                            "average_rating": (
                                round(avg_rating, 2) if avg_rating else None
                            ),
                            "avg_rating": (
                                round(avg_rating, 2) if avg_rating else None
                            ),  # Keep both for compatibility
                            "recent_prompts": total_prompts,  # For now, use total as recent count
                        },
                    }
                )

        except Exception as e:
            self.logger.error(f"Stats error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to get statistics: {str(e)}"},
                status=500,
            )

    async def get_settings(self, request):
        """Get current settings."""
        try:
            from .config import PromptManagerConfig
            
            # Return actual configuration settings
            return web.json_response({
                "success": True, 
                "settings": {
                    "result_timeout": PromptManagerConfig.RESULT_TIMEOUT,
                    "webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE
                }
            })
        except Exception as e:
            return web.json_response(
                {"success": False, "error": f"Failed to get settings: {str(e)}"},
                status=500,
            )

    async def save_settings(self, request):
        """Save settings."""
        try:
            data = await request.json()
            # For now, just acknowledge the save
            # In the future, we could store this in database or config file
            return web.json_response(
                {"success": True, "message": "Settings saved successfully"}
            )
        except Exception as e:
            return web.json_response(
                {"success": False, "error": f"Failed to save settings: {str(e)}"},
                status=500,
            )

    async def update_prompt(self, request):
        """Update prompt text."""
        try:
            prompt_id = int(request.match_info["prompt_id"])
            data = await request.json()
            new_text = data.get("text", "").strip()

            if not new_text:
                return web.json_response(
                    {"success": False, "error": "Text cannot be empty"}, status=400
                )

            # Update the prompt in database
            with self.db.model.get_connection() as conn:
                cursor = conn.execute(
                    "UPDATE prompts SET text = ?, updated_at = ? WHERE id = ?",
                    (
                        new_text,
                        datetime.datetime.now(datetime.timezone.utc).isoformat(),
                        prompt_id,
                    ),
                )
                conn.commit()

                if cursor.rowcount > 0:
                    return web.json_response(
                        {"success": True, "message": "Prompt updated successfully"}
                    )
                else:
                    return web.json_response(
                        {"success": False, "error": "Prompt not found"}, status=404
                    )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Update prompt error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to update prompt: {str(e)}"},
                status=500,
            )

    async def update_prompt_rating(self, request):
        """Update prompt rating."""
        try:
            prompt_id = int(request.match_info["prompt_id"])
            data = await request.json()
            rating = data.get("rating")

            if rating is not None and (rating < 1 or rating > 5):
                return web.json_response(
                    {"success": False, "error": "Rating must be between 1 and 5"},
                    status=400,
                )

            with self.db.model.get_connection() as conn:
                cursor = conn.execute(
                    "UPDATE prompts SET rating = ?, updated_at = ? WHERE id = ?",
                    (
                        rating,
                        datetime.datetime.now(datetime.timezone.utc).isoformat(),
                        prompt_id,
                    ),
                )
                conn.commit()

                if cursor.rowcount > 0:
                    return web.json_response(
                        {"success": True, "message": "Rating updated successfully"}
                    )
                else:
                    return web.json_response(
                        {"success": False, "error": "Prompt not found"}, status=404
                    )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Update rating error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to update rating: {str(e)}"},
                status=500,
            )

    async def add_prompt_tag(self, request):
        """Add tag to prompt."""
        try:
            prompt_id = int(request.match_info["prompt_id"])
            data = await request.json()
            new_tag = data.get("tag", "").strip()

            if not new_tag:
                return web.json_response(
                    {"success": False, "error": "Tag cannot be empty"}, status=400
                )

            # Get current prompt
            prompt = self.db.get_prompt_by_id(prompt_id)
            if not prompt:
                return web.json_response(
                    {"success": False, "error": "Prompt not found"}, status=404
                )

            # Get current tags
            current_tags = prompt.get("tags", [])
            if not isinstance(current_tags, list):
                current_tags = []

            # Add new tag if not already present
            if new_tag not in current_tags:
                current_tags.append(new_tag)

                # Update database
                with self.db.model.get_connection() as conn:
                    cursor = conn.execute(
                        "UPDATE prompts SET tags = ?, updated_at = ? WHERE id = ?",
                        (
                            json.dumps(current_tags),
                            datetime.datetime.now(datetime.timezone.utc).isoformat(),
                            prompt_id,
                        ),
                    )
                    conn.commit()

            return web.json_response(
                {"success": True, "message": "Tag added successfully"}
            )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Add tag error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to add tag: {str(e)}"}, status=500
            )

    async def add_tags_to_prompt(self, request):
        """Add multiple tags to a single prompt."""
        try:
            data = await request.json()
            prompt_id = data.get("prompt_id")
            new_tags = data.get("tags", [])

            if not prompt_id:
                return web.json_response(
                    {"success": False, "error": "Prompt ID is required"}, status=400
                )

            if not new_tags or not isinstance(new_tags, list):
                return web.json_response(
                    {"success": False, "error": "Tags must be a non-empty list"}, status=400
                )

            # Get current prompt
            prompt = self.db.get_prompt_by_id(prompt_id)
            if not prompt:
                return web.json_response(
                    {"success": False, "error": "Prompt not found"}, status=404
                )

            # Get current tags
            current_tags = prompt.get("tags", [])
            if not isinstance(current_tags, list):
                current_tags = []

            # Add new tags if not already present
            tags_added = 0
            for new_tag in new_tags:
                new_tag = new_tag.strip()
                if new_tag and new_tag not in current_tags:
                    current_tags.append(new_tag)
                    tags_added += 1

            # Update database if any tags were added
            if tags_added > 0:
                with self.db.model.get_connection() as conn:
                    cursor = conn.execute(
                        "UPDATE prompts SET tags = ?, updated_at = ? WHERE id = ?",
                        (
                            json.dumps(current_tags),
                            datetime.datetime.now(datetime.timezone.utc).isoformat(),
                            prompt_id,
                        ),
                    )
                    conn.commit()

            message = f"{tags_added} tag(s) added successfully"
            if tags_added == 0:
                message = "No new tags to add (all tags already exist)"

            return web.json_response(
                {"success": True, "message": message, "tags_added": tags_added}
            )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Add tags error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500
            )

    async def remove_prompt_tag(self, request):
        """Remove tag from prompt."""
        try:
            prompt_id = int(request.match_info["prompt_id"])
            data = await request.json()
            tag_to_remove = data.get("tag", "").strip()

            # Get current prompt
            prompt = self.db.get_prompt_by_id(prompt_id)
            if not prompt:
                return web.json_response(
                    {"success": False, "error": "Prompt not found"}, status=404
                )

            # Get current tags
            current_tags = prompt.get("tags", [])
            if not isinstance(current_tags, list):
                current_tags = []

            # Remove tag if present
            if tag_to_remove in current_tags:
                current_tags.remove(tag_to_remove)

                # Update database
                with self.db.model.get_connection() as conn:
                    cursor = conn.execute(
                        "UPDATE prompts SET tags = ?, updated_at = ? WHERE id = ?",
                        (
                            json.dumps(current_tags),
                            datetime.datetime.now(datetime.timezone.utc).isoformat(),
                            prompt_id,
                        ),
                    )
                    conn.commit()

            return web.json_response(
                {"success": True, "message": "Tag removed successfully"}
            )

        except ValueError:
            return web.json_response(
                {"success": False, "error": "Invalid prompt ID"}, status=400
            )
        except Exception as e:
            self.logger.error(f"Remove tag error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to remove tag: {str(e)}"},
                status=500,
            )

    async def bulk_delete_prompts(self, request):
        """Bulk delete prompts."""
        try:
            data = await request.json()
            prompt_ids = data.get("prompt_ids", [])

            if not prompt_ids:
                return web.json_response(
                    {"success": False, "error": "No prompt IDs provided"}, status=400
                )

            deleted_count = 0
            with self.db.model.get_connection() as conn:
                for prompt_id in prompt_ids:
                    # First delete related images to avoid foreign key constraint
                    conn.execute("DELETE FROM generated_images WHERE prompt_id = ?", (prompt_id,))
                    # Then delete the prompt
                    cursor = conn.execute(
                        "DELETE FROM prompts WHERE id = ?", (prompt_id,)
                    )
                    if cursor.rowcount > 0:
                        deleted_count += 1
                conn.commit()

            return web.json_response(
                {
                    "success": True,
                    "message": f"Deleted {deleted_count} prompts",
                    "deleted_count": deleted_count,
                }
            )

        except Exception as e:
            self.logger.error(f"Bulk delete error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to delete prompts: {str(e)}"},
                status=500,
            )

    async def bulk_add_tags(self, request):
        """Bulk add tags to prompts."""
        try:
            data = await request.json()
            prompt_ids = data.get("prompt_ids", [])
            new_tags = data.get("tags", [])

            if not prompt_ids or not new_tags:
                return web.json_response(
                    {"success": False, "error": "No prompt IDs or tags provided"},
                    status=400,
                )

            updated_count = 0
            with self.db.model.get_connection() as conn:
                for prompt_id in prompt_ids:
                    # Get current tags
                    cursor = conn.execute(
                        "SELECT tags FROM prompts WHERE id = ?", (prompt_id,)
                    )
                    row = cursor.fetchone()
                    if row:
                        current_tags = []
                        if row["tags"]:
                            try:
                                current_tags = json.loads(row["tags"])
                                if not isinstance(current_tags, list):
                                    current_tags = []
                            except:
                                current_tags = []

                        # Add new tags
                        for tag in new_tags:
                            if tag not in current_tags:
                                current_tags.append(tag)

                        # Update database
                        cursor = conn.execute(
                            "UPDATE prompts SET tags = ?, updated_at = ? WHERE id = ?",
                            (
                                json.dumps(current_tags),
                                datetime.datetime.now(
                                    datetime.timezone.utc
                                ).isoformat(),
                                prompt_id,
                            ),
                        )
                        if cursor.rowcount > 0:
                            updated_count += 1

                conn.commit()

            return web.json_response(
                {
                    "success": True,
                    "message": f"Added tags to {updated_count} prompts",
                    "updated_count": updated_count,
                }
            )

        except Exception as e:
            self.logger.error(f"Bulk add tags error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500
            )

    async def bulk_set_category(self, request):
        """Bulk set category for prompts."""
        try:
            data = await request.json()
            prompt_ids = data.get("prompt_ids", [])
            category = data.get("category", "").strip()

            if not prompt_ids:
                return web.json_response(
                    {"success": False, "error": "No prompt IDs provided"}, status=400
                )

            updated_count = 0
            with self.db.model.get_connection() as conn:
                for prompt_id in prompt_ids:
                    cursor = conn.execute(
                        "UPDATE prompts SET category = ?, updated_at = ? WHERE id = ?",
                        (
                            category,
                            datetime.datetime.now(datetime.timezone.utc).isoformat(),
                            prompt_id,
                        ),
                    )
                    if cursor.rowcount > 0:
                        updated_count += 1
                conn.commit()

            return web.json_response(
                {
                    "success": True,
                    "message": f"Set category for {updated_count} prompts",
                    "updated_count": updated_count,
                }
            )

        except Exception as e:
            self.logger.error(f"Bulk set category error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to set category: {str(e)}"},
                status=500,
            )

    async def export_prompts(self, request):
        """Export all prompts to JSON."""
        try:
            # Get all prompts
            prompts = self.db.search_prompts(limit=10000)

            # Create export data
            export_data = {
                "export_date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                "total_prompts": len(prompts),
                "prompts": prompts,
            }

            # Return as JSON download
            json_data = json.dumps(export_data, indent=2, ensure_ascii=False)

            return web.Response(
                text=json_data,
                content_type="application/json",
                headers={
                    "Content-Disposition": f'attachment; filename="prompt_manager_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}.json"'
                },
            )

        except Exception as e:
            self.logger.error(f"Export error: {e}")
            return web.json_response(
                {"success": False, "error": f"Failed to export prompts: {str(e)}"},
                status=500,
            )

    # Gallery-related endpoints
    def _clean_nan_recursive(self, obj):
        """Recursively clean NaN values from nested data structures."""
        if isinstance(obj, dict):
            return {key: self._clean_nan_recursive(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [self._clean_nan_recursive(item) for item in obj]
        elif isinstance(obj, float) and str(obj) == 'nan':
            return None
        else:
            return obj

    async def get_prompt_images(self, request):
        """Get all images for a specific prompt."""
        try:
            prompt_id = request.match_info["prompt_id"]
            images = self.db.get_prompt_images(prompt_id)
            
            # Clean up any NaN values that cause JSON parsing errors (recursive)
            cleaned_images = [self._clean_nan_recursive(image) for image in images]
            
            # Additional fallback: convert to JSON string and clean NaN values manually
            import json
            import re
            try:
                response_data = {
                    'success': True,
                    'images': cleaned_images
                }
                # Convert to JSON string
                json_str = json.dumps(response_data, default=str)
                
                # Clean any remaining NaN values with regex
                json_str = re.sub(r':\s*NaN', ': null', json_str)
                json_str = re.sub(r'\[\s*NaN\s*\]', '[null]', json_str)
                json_str = re.sub(r',\s*NaN\s*,', ', null,', json_str)
                json_str = re.sub(r',\s*NaN\s*\]', ', null]', json_str)
                json_str = re.sub(r'\[\s*NaN\s*,', '[null,', json_str)
                
                # Parse back to verify it's valid JSON
                cleaned_data = json.loads(json_str)
                
                return web.json_response(cleaned_data)
            except Exception as json_error:
                self.logger.error(f"JSON cleaning error: {json_error}")
                # Fallback to original response
                return web.json_response({
                    'success': True,
                    'images': cleaned_images
                })
            
        except Exception as e:
            self.logger.error(f"Get prompt images error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def get_recent_images(self, request):
        """Get recently generated images."""
        try:
            limit = int(request.query.get('limit', 50))
            images = self.db.get_recent_images(limit)
            
            return web.json_response({
                'success': True,
                'images': images
            })
        except Exception as e:
            self.logger.error(f"Get recent images error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def search_images(self, request):
        """Search images by prompt text."""
        try:
            query = request.query.get('q', '')
            if not query:
                return web.json_response({
                    'success': False,
                    'error': 'Search query required'
                }, status=400)
            
            images = self.db.search_images_by_prompt(query)
            
            return web.json_response({
                'success': True,
                'images': images,
                'query': query
            })
        except Exception as e:
            self.logger.error(f"Search images error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def serve_image(self, request):
        """Serve the actual image file."""
        try:
            image_id = int(request.match_info["image_id"])
            image = self.db.get_image_by_id(image_id)
            
            if not image:
                return web.json_response({'error': 'Image not found'}, status=404)
            
            import os
            from pathlib import Path
            
            image_path = Path(image['image_path'])
            if not image_path.exists():
                return web.json_response({'error': 'Image file not found'}, status=404)
            
            # Determine content type based on file extension
            content_type = 'image/jpeg'
            if image_path.suffix.lower() in ['.png']:
                content_type = 'image/png'
            elif image_path.suffix.lower() in ['.webp']:
                content_type = 'image/webp'
            elif image_path.suffix.lower() in ['.gif']:
                content_type = 'image/gif'
            
            # Read and serve the file
            with open(image_path, 'rb') as f:
                file_data = f.read()
            
            return web.Response(
                body=file_data,
                content_type=content_type,
                headers={
                    'Cache-Control': 'public, max-age=3600',
                    'Content-Length': str(len(file_data))
                }
            )
            
        except ValueError:
            return web.json_response({'error': 'Invalid image ID'}, status=400)
        except Exception as e:
            self.logger.error(f"Serve image error: {e}")
            return web.json_response({'error': str(e)}, status=500)

    async def link_image_to_prompt(self, request):
        """Link a generated image to a prompt."""
        try:
            data = await request.json()
            prompt_id = data.get('prompt_id')
            image_path = data.get('image_path')
            metadata = data.get('metadata', {})
            
            if not prompt_id or not image_path:
                return web.json_response({
                    'success': False,
                    'error': 'prompt_id and image_path are required'
                }, status=400)
            
            # Check if image file exists
            import os
            if not os.path.exists(image_path):
                return web.json_response({
                    'success': False,
                    'error': 'Image file not found'
                }, status=404)
            
            # Link image to prompt
            image_id = self.db.link_image_to_prompt(prompt_id, image_path, metadata)
            
            return web.json_response({
                'success': True,
                'image_id': image_id,
                'message': 'Image linked successfully'
            })
            
        except Exception as e:
            self.logger.error(f"Link image error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def delete_image(self, request):
        """Delete an image record."""
        try:
            image_id = int(request.match_info["image_id"])
            success = self.db.delete_image(image_id)
            
            if success:
                return web.json_response({
                    'success': True,
                    'message': 'Image deleted successfully'
                })
            else:
                return web.json_response({
                    'success': False,
                    'error': 'Image not found'
                }, status=404)
                
        except ValueError:
            return web.json_response({'error': 'Invalid image ID'}, status=400)
        except Exception as e:
            self.logger.error(f"Delete image error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    # Diagnostic endpoints
    async def run_diagnostics(self, request):
        """Run system diagnostics."""
        try:
            # Simple diagnostics without importing complex modules
            import os
            import sqlite3
            
            results = {}
            
            # Check database
            try:
                db_path = "prompts.db"
                if os.path.exists(db_path):
                    with sqlite3.connect(db_path) as conn:
                        conn.row_factory = sqlite3.Row
                        cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
                        prompt_count = cursor.fetchone()['count']
                        
                        # Check if images table exists
                        cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='generated_images'")
                        has_images_table = cursor.fetchone() is not None
                        
                        if has_images_table:
                            cursor = conn.execute("SELECT COUNT(*) as count FROM generated_images")
                            image_count = cursor.fetchone()['count']
                        else:
                            image_count = 0
                        
                        results['database'] = {
                            'status': 'ok',
                            'prompt_count': prompt_count,
                            'has_images_table': has_images_table,
                            'image_count': image_count
                        }
                else:
                    results['database'] = {
                        'status': 'error',
                        'message': f'Database file not found: {db_path}'
                    }
            except Exception as e:
                results['database'] = {
                    'status': 'error',
                    'message': f'Database error: {str(e)}'
                }
            
            # Check dependencies
            dependencies = {}
            try:
                import watchdog
                dependencies['watchdog'] = True
            except ImportError:
                dependencies['watchdog'] = False
            
            try:
                from PIL import Image
                dependencies['PIL'] = True
            except ImportError:
                dependencies['PIL'] = False
            
            try:
                import sqlite3
                dependencies['sqlite3'] = True
            except ImportError:
                dependencies['sqlite3'] = False
            
            results['dependencies'] = {
                'status': 'ok' if all(dependencies.values()) else 'error',
                'dependencies': dependencies
            }
            
            # Check output directories
            output_dirs = []
            potential_dirs = ["output", "../output", "../../output"]
            
            for dir_path in potential_dirs:
                abs_path = os.path.abspath(dir_path)
                if os.path.exists(abs_path):
                    output_dirs.append(abs_path)
            
            results['comfyui_output'] = {
                'status': 'ok' if output_dirs else 'warning',
                'output_dirs': output_dirs
            }
            
            return web.json_response({
                'success': True,
                'diagnostics': results
            })
            
        except Exception as e:
            self.logger.error(f"Diagnostics error: {e}", exc_info=True)
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def test_image_link(self, request):
        """Test creating an image link."""
        try:
            data = await request.json()
            prompt_id = data.get('prompt_id')
            test_image_path = data.get('image_path', '/test/fake/image.png')
            
            if not prompt_id:
                return web.json_response({
                    'success': False,
                    'error': 'prompt_id is required'
                }, status=400)
            
            # Test linking directly using the database manager
            test_metadata = {
                'file_info': {
                    'size': 1024000,
                    'dimensions': [512, 512],
                    'format': 'PNG'
                },
                'workflow': {'test': True},
                'prompt': {'test_prompt': 'This is a test image'}
            }
            
            try:
                image_id = self.db.link_image_to_prompt(
                    prompt_id=str(prompt_id),
                    image_path=test_image_path,
                    metadata=test_metadata
                )
                
                return web.json_response({
                    'success': True,
                    'result': {
                        'status': 'ok',
                        'image_id': image_id,
                        'message': f'Test image linked successfully with ID {image_id}'
                    }
                })
            except Exception as e:
                return web.json_response({
                    'success': False,
                    'result': {
                        'status': 'error',
                        'message': f'Failed to create test link: {str(e)}'
                    }
                })
                
        except Exception as e:
            self.logger.error(f"Test link error: {e}", exc_info=True)
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def run_maintenance(self, request):
        """
        Comprehensive database maintenance endpoint.
        POST /prompt_manager/maintenance
        """
        try:
            data = await request.json() if request.content_type == 'application/json' else {}
            operations = data.get('operations', ['cleanup_duplicates', 'vacuum', 'cleanup_orphaned_images'])
            
            results = {}
            
            # Clean up duplicate prompts
            if 'cleanup_duplicates' in operations:
                try:
                    duplicates_removed = self.db.cleanup_duplicates()
                    results['cleanup_duplicates'] = {
                        'success': True,
                        'removed_count': duplicates_removed,
                        'message': f'Removed {duplicates_removed} duplicate prompts'
                    }
                except Exception as e:
                    results['cleanup_duplicates'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to cleanup duplicates'
                    }
            
            # Vacuum database
            if 'vacuum' in operations:
                try:
                    self.db.model.vacuum_database()
                    results['vacuum'] = {
                        'success': True,
                        'message': 'Database vacuum completed successfully'
                    }
                except Exception as e:
                    results['vacuum'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to vacuum database'
                    }
            
            # Clean up orphaned image records
            if 'cleanup_orphaned_images' in operations:
                try:
                    orphaned_removed = self.db.cleanup_missing_images()
                    results['cleanup_orphaned_images'] = {
                        'success': True,
                        'removed_count': orphaned_removed,
                        'message': f'Removed {orphaned_removed} orphaned image records'
                    }
                except Exception as e:
                    results['cleanup_orphaned_images'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to cleanup orphaned images'
                    }
            
            # Check for potential duplicate hashes
            if 'check_hash_duplicates' in operations:
                try:
                    with self.db.model.get_connection() as conn:
                        cursor = conn.execute("""
                            SELECT hash, COUNT(*) as count 
                            FROM prompts 
                            WHERE hash IS NOT NULL 
                            GROUP BY hash 
                            HAVING COUNT(*) > 1
                        """)
                        hash_duplicates = cursor.fetchall()
                        
                        results['check_hash_duplicates'] = {
                            'success': True,
                            'duplicate_hashes': len(hash_duplicates),
                            'message': f'Found {len(hash_duplicates)} duplicate hash groups'
                        }
                except Exception as e:
                    results['check_hash_duplicates'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to check hash duplicates'
                    }
            
            # Get database statistics
            if 'statistics' in operations:
                try:
                    db_info = self.db.model.get_database_info()
                    results['statistics'] = {
                        'success': True,
                        'info': db_info,
                        'message': 'Database statistics retrieved'
                    }
                except Exception as e:
                    results['statistics'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to get database statistics'
                    }
            
            # Prune orphaned prompts (prompts with no linked images)
            if 'prune_orphaned_prompts' in operations:
                try:
                    with self.db.model.get_connection() as conn:
                        # Find prompts that have no images linked to them
                        cursor = conn.execute("""
                            SELECT p.id 
                            FROM prompts p 
                            LEFT JOIN generated_images gi ON p.id = gi.prompt_id 
                            WHERE gi.prompt_id IS NULL
                        """)
                        orphaned_prompts = cursor.fetchall()
                        orphaned_count = len(orphaned_prompts)
                        
                        if orphaned_count > 0:
                            # Delete the orphaned prompts
                            orphaned_ids = [row['id'] for row in orphaned_prompts]
                            placeholders = ','.join(['?'] * len(orphaned_ids))
                            cursor = conn.execute(f"DELETE FROM prompts WHERE id IN ({placeholders})", orphaned_ids)
                            conn.commit()
                            removed_count = cursor.rowcount
                        else:
                            removed_count = 0
                    
                    results['prune_orphaned_prompts'] = {
                        'success': True,
                        'removed_count': removed_count,
                        'message': f'Removed {removed_count} orphaned prompts (prompts with no linked images)'
                    }
                except Exception as e:
                    results['prune_orphaned_prompts'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to prune orphaned prompts'
                    }
            
            # Check for consistency issues
            if 'check_consistency' in operations:
                try:
                    consistency_issues = []
                    
                    with self.db.model.get_connection() as conn:
                        # Check for prompts with invalid JSON in tags
                        cursor = conn.execute("SELECT id, tags FROM prompts WHERE tags IS NOT NULL")
                        for row in cursor.fetchall():
                            try:
                                if row['tags']:
                                    json.loads(row['tags'])
                            except json.JSONDecodeError:
                                consistency_issues.append(f"Prompt {row['id']} has invalid JSON in tags")
                        
                        # Check for orphaned foreign key references
                        cursor = conn.execute("""
                            SELECT gi.id, gi.prompt_id 
                            FROM generated_images gi 
                            LEFT JOIN prompts p ON gi.prompt_id = p.id 
                            WHERE p.id IS NULL
                        """)
                        orphaned_refs = cursor.fetchall()
                        for ref in orphaned_refs:
                            consistency_issues.append(f"Image {ref['id']} references non-existent prompt {ref['prompt_id']}")
                    
                    results['check_consistency'] = {
                        'success': True,
                        'issues_found': len(consistency_issues),
                        'issues': consistency_issues[:10],  # Limit to first 10 issues
                        'message': f'Found {len(consistency_issues)} consistency issues'
                    }
                except Exception as e:
                    results['check_consistency'] = {
                        'success': False,
                        'error': str(e),
                        'message': 'Failed to check database consistency'
                    }
            
            # Overall success status
            all_successful = all(result.get('success', False) for result in results.values())
            
            return web.json_response({
                'success': True,
                'operations_completed': len(results),
                'all_successful': all_successful,
                'results': results,
                'message': f'Maintenance completed: {len(results)} operations processed'
            })
            
        except Exception as e:
            self.logger.error(f"Maintenance error: {e}", exc_info=True)
            return web.json_response({
                'success': False,
                'error': f'Maintenance failed: {str(e)}'
            }, status=500)

    async def backup_database(self, request):
        """
        Backup the entire prompts.db database file.
        GET /prompt_manager/backup
        """
        try:
            import os
            import shutil
            import tempfile
            from pathlib import Path

            # Get the database path
            db_path = "prompts.db"
            
            if not os.path.exists(db_path):
                return web.json_response({
                    'success': False,
                    'error': 'Database file not found'
                }, status=404)

            # Create a temporary copy of the database
            with tempfile.NamedTemporaryFile(delete=False, suffix='.db') as temp_file:
                temp_path = temp_file.name
                
            # Copy the database file
            shutil.copy2(db_path, temp_path)
            
            # Read the file content
            with open(temp_path, 'rb') as f:
                file_data = f.read()
            
            # Clean up temporary file
            os.unlink(temp_path)
            
            # Generate filename with timestamp
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"prompts_backup_{timestamp}.db"
            
            return web.Response(
                body=file_data,
                content_type='application/octet-stream',
                headers={
                    'Content-Disposition': f'attachment; filename="{filename}"',
                    'Content-Length': str(len(file_data))
                }
            )

        except Exception as e:
            self.logger.error(f"Backup error: {e}", exc_info=True)
            return web.json_response({
                'success': False,
                'error': f'Failed to backup database: {str(e)}'
            }, status=500)

    async def restore_database(self, request):
        """
        Restore the prompts.db database from uploaded file.
        POST /prompt_manager/restore
        """
        try:
            import os
            import shutil
            import tempfile
            import sqlite3
            from pathlib import Path

            # Get the uploaded file
            reader = await request.multipart()
            field = await reader.next()
            
            if not field or field.name != 'database_file':
                return web.json_response({
                    'success': False,
                    'error': 'No database file uploaded. Expected field name: database_file'
                }, status=400)

            # Read the uploaded file content
            file_data = await field.read()
            
            if not file_data:
                return web.json_response({
                    'success': False,
                    'error': 'Uploaded file is empty'
                }, status=400)

            # Create a temporary file to validate the database
            with tempfile.NamedTemporaryFile(delete=False, suffix='.db') as temp_file:
                temp_path = temp_file.name
                temp_file.write(file_data)

            try:
                # Validate that it's a valid SQLite database with expected structure
                with sqlite3.connect(temp_path) as conn:
                    conn.row_factory = sqlite3.Row
                    
                    # Check if it has the prompts table
                    cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'")
                    if not cursor.fetchone():
                        raise ValueError("Database does not contain a 'prompts' table")
                    
                    # Check basic structure of prompts table
                    cursor = conn.execute("PRAGMA table_info(prompts)")
                    columns = [row['name'] for row in cursor.fetchall()]
                    required_columns = ['id', 'text', 'created_at']
                    
                    for col in required_columns:
                        if col not in columns:
                            raise ValueError(f"Database missing required column: {col}")
                    
                    # Get basic stats for validation
                    cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
                    prompt_count = cursor.fetchone()['count']

                # If validation passes, backup current database and restore
                db_path = "prompts.db"
                backup_path = f"{db_path}.backup_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}"
                
                # Create backup of current database if it exists
                if os.path.exists(db_path):
                    shutil.copy2(db_path, backup_path)
                    self.logger.info(f"Current database backed up to: {backup_path}")

                # Replace current database with uploaded one
                shutil.copy2(temp_path, db_path)
                
                # Reinitialize the database connection
                self.db = PromptDatabase()
                
                return web.json_response({
                    'success': True,
                    'message': f'Database restored successfully. Found {prompt_count} prompts.',
                    'prompt_count': prompt_count,
                    'backup_created': backup_path if os.path.exists(db_path) else None
                })

            except sqlite3.Error as e:
                return web.json_response({
                    'success': False,
                    'error': f'Invalid SQLite database: {str(e)}'
                }, status=400)
            except ValueError as e:
                return web.json_response({
                    'success': False,
                    'error': str(e)
                }, status=400)
            finally:
                # Clean up temporary file
                if os.path.exists(temp_path):
                    os.unlink(temp_path)

        except Exception as e:
            self.logger.error(f"Restore error: {e}", exc_info=True)
            return web.json_response({
                'success': False,
                'error': f'Failed to restore database: {str(e)}'
            }, status=500)

    async def scan_images(self, request):
        """
        Scan ComfyUI output images for prompt metadata and add them to the database.
        Streams progress updates to the client.
        """
        import json
        import asyncio
        from pathlib import Path
        from aiohttp import web
        
        async def stream_response():
            try:
                self.logger.info("Starting image scan operation")
                
                # Clear any active prompt tracking timers to avoid conflicts
                # Note: For now we'll skip this since we don't have easy access to the tracker instance
                # This is mainly important if scans are run during active generation, which is rare
                self.logger.info("Starting scan (timer clearing not implemented yet)")
                
                # Find ComfyUI output directory
                output_dir = self._find_comfyui_output_dir()
                if not output_dir:
                    self.logger.error("ComfyUI output directory not found")
                    yield f"data: {json.dumps({'type': 'error', 'message': 'ComfyUI output directory not found'})}\n\n"
                    return
                
                yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Scanning for PNG files...', 'processed': 0, 'found': 0})}\n\n"
                
                # Find all PNG files
                png_files = list(Path(output_dir).rglob("*.png"))
                total_files = len(png_files)
                
                if total_files == 0:
                    yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'found': 0, 'added': 0})}\n\n"
                    return
                
                yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': f'Found {total_files} PNG files to process...', 'processed': 0, 'found': 0})}\n\n"
                
                processed_count = 0
                found_count = 0
                added_count = 0
                linked_count = 0  # Count of images linked to existing prompts
                
                for i, png_file in enumerate(png_files):
                    try:
                        # Extract metadata from PNG
                        metadata = self._extract_comfyui_metadata(str(png_file))
                        processed_count += 1
                        
                        if metadata:
                            self.logger.debug(f"Found metadata in {os.path.basename(png_file)}: {list(metadata.keys())}")
                            
                            # Parse ComfyUI prompt data
                            parsed_data = self._parse_comfyui_prompt(metadata)
                            self.logger.debug(f"Parsed data keys: {list(parsed_data.keys())}, has prompt: {bool(parsed_data.get('prompt'))}, has parameters: {bool(parsed_data.get('parameters'))}")
                            
                            # Check if we found any meaningful prompt data
                            if parsed_data.get('prompt') or parsed_data.get('parameters'):
                                found_count += 1
                                
                                # Extract readable prompt text
                                prompt_text = self._extract_readable_prompt(parsed_data)
                                
                                # Debug: print what we found and its type
                                if prompt_text:
                                    self.logger.debug(f"Found prompt in {os.path.basename(png_file)} (type: {type(prompt_text)}): {str(prompt_text)[:100]}...")
                                else:
                                    self.logger.debug(f"No readable prompt found in {os.path.basename(png_file)}, parsed_data keys: {list(parsed_data.keys())}")
                                
                                # Ensure prompt_text is a string
                                if prompt_text and not isinstance(prompt_text, str):
                                    self.logger.debug(f"Converting prompt_text from {type(prompt_text)} to string")
                                    prompt_text = str(prompt_text)
                                
                                if prompt_text and prompt_text.strip():
                                    # Try to save to database (will skip duplicates)
                                    try:
                                        # Generate hash for duplicate detection
                                        try:
                                            from ..utils.hashing import generate_prompt_hash
                                        except ImportError:
                                            import sys
                                            current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                                            sys.path.insert(0, current_dir)
                                            from utils.hashing import generate_prompt_hash
                                        
                                        prompt_hash = generate_prompt_hash(prompt_text.strip())
                                        self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
                                        
                                        # Check if prompt already exists
                                        existing = self.db.get_prompt_by_hash(prompt_hash)
                                        if existing:
                                            self.logger.debug(f"Found existing prompt ID {existing['id']} for image {os.path.basename(png_file)}")
                                            # Link image to existing prompt
                                            try:
                                                self.db.link_image_to_prompt(existing['id'], str(png_file))
                                                linked_count += 1
                                                self.logger.debug(f"Linked image {os.path.basename(png_file)} to existing prompt {existing['id']}")
                                            except Exception as e:
                                                self.logger.error(f"Failed to link image {png_file} to existing prompt: {e}")
                                        else:
                                            # Save new prompt
                                            self.logger.debug(f"Saving new prompt from {os.path.basename(png_file)}")
                                            prompt_id = self.db.save_prompt(
                                                text=prompt_text.strip(),
                                                category='scanned',
                                                tags=['auto-scanned'],
                                                notes=f'Auto-scanned from {os.path.basename(png_file)}',
                                                prompt_hash=prompt_hash
                                            )
                                            
                                            if prompt_id:
                                                added_count += 1
                                                self.logger.info(f"Successfully saved new prompt with ID {prompt_id} from {os.path.basename(png_file)}")
                                                
                                                # Link image to new prompt
                                                try:
                                                    self.db.link_image_to_prompt(prompt_id, str(png_file))
                                                    self.logger.debug(f"Linked image {os.path.basename(png_file)} to new prompt {prompt_id}")
                                                except Exception as e:
                                                    self.logger.error(f"Failed to link image {png_file} to new prompt: {e}")
                                            else:
                                                self.logger.error(f"Failed to save prompt from {os.path.basename(png_file)} - no ID returned")
                                        
                                    except Exception as e:
                                        self.logger.error(f"Failed to save prompt from {png_file}: {e}")
                        
                        # Update progress every 10 files or so
                        if i % 10 == 0 or i == total_files - 1:
                            progress = int((i + 1) / total_files * 100)
                            status = f"Processing file {i + 1}/{total_files}..."
                            
                            yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': status, 'processed': processed_count, 'found': found_count})}\n\n"
                            
                            # Small delay to allow UI updates
                            await asyncio.sleep(0.01)
                    
                    except Exception as e:
                        self.logger.error(f"Error processing {png_file}: {e}")
                        continue
                
                # Send completion message
                self.logger.info(f"Scan completed: processed={processed_count}, found={found_count}, new_prompts_added={added_count}, images_linked_to_existing={linked_count}")
                yield f"data: {json.dumps({'type': 'complete', 'processed': processed_count, 'found': found_count, 'added': added_count, 'linked': linked_count})}\n\n"
                
            except Exception as e:
                self.logger.error(f"Scan error: {e}")
                import traceback
                self.logger.error(f"Scan error traceback: {traceback.format_exc()}")
                yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
        
        # Return streaming response
        response = web.StreamResponse(
            status=200,
            reason='OK',
            headers={
                'Content-Type': 'text/event-stream',
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive'
            }
        )
        
        await response.prepare(request)
        
        async for chunk in stream_response():
            await response.write(chunk.encode('utf-8'))
        
        await response.write_eof()
        return response
    
    def _find_comfyui_output_dir(self):
        """Find the ComfyUI output directory."""
        
        # Common ComfyUI installation paths
        possible_paths = [
            os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "..", "..", "output"),
            os.path.join(os.path.expanduser("~"), "ComfyUI", "output"),
            os.path.join(os.getcwd(), "output"),
            os.path.join(os.getcwd(), "..", "output"),
            os.path.join(os.getcwd(), "..", "..", "output"),
        ]
        
        for path in possible_paths:
            abs_path = os.path.abspath(path)
            if os.path.exists(abs_path) and os.path.isdir(abs_path):
                self.logger.info(f"Found ComfyUI output directory: {abs_path}")
                return abs_path
        
        self.logger.warning("ComfyUI output directory not found")
        return None
    
    def _extract_comfyui_metadata(self, image_path):
        """Extract ComfyUI metadata from a PNG file."""
        try:
            with Image.open(image_path) as img:
                metadata = {}
                if hasattr(img, 'text'):
                    for key, value in img.text.items():
                        metadata[key] = value
                return metadata
        except Exception as e:
            self.logger.error(f"Error reading {image_path}: {e}")
            return {}
    
    def _parse_comfyui_prompt(self, metadata):
        """Parse ComfyUI and A1111 prompt data from metadata."""
        result = {
            'prompt': None,
            'workflow': None,
            'parameters': {},
            'positive_prompt': None,
            'negative_prompt': None
        }
        
        # Check for A1111 style parameters first (like parse-metadata.py)
        if "parameters" in metadata:
            params = metadata["parameters"]
            lines = params.splitlines()
            if lines:
                result['positive_prompt'] = lines[0].strip()
            for line in lines:
                if line.lower().startswith("negative prompt:"):
                    result['negative_prompt'] = line.split(":", 1)[1].strip()
                    break
            # Store raw parameters too
            result['parameters']['parameters'] = params
        
        # If no A1111 format found, proceed with ComfyUI parsing
        if result['positive_prompt'] is None:
            # Check for direct prompt field
            if 'prompt' in metadata:
                try:
                    prompt_data = json.loads(metadata['prompt'])
                    result['prompt'] = prompt_data
                except json.JSONDecodeError:
                    result['prompt'] = metadata['prompt']
            
            # Check for workflow
            if 'workflow' in metadata:
                try:
                    workflow_data = json.loads(metadata['workflow'])
                    result['workflow'] = workflow_data
                except json.JSONDecodeError:
                    result['workflow'] = metadata['workflow']
            
            # Check for other common ComfyUI fields
            common_fields = ['positive', 'negative', 'steps', 'cfg', 'sampler', 'scheduler', 'seed']
            for field in common_fields:
                if field in metadata:
                    try:
                        result['parameters'][field] = json.loads(metadata[field])
                    except json.JSONDecodeError:
                        result['parameters'][field] = metadata[field]
        
        return result
    
    def _extract_readable_prompt(self, parsed_data):
        """Extract human-readable prompt text from ComfyUI/A1111 data using improved logic."""
        import json
        
        # Helper function to convert any value to string safely
        def safe_to_string(value):
            if isinstance(value, str):
                return value
            elif isinstance(value, list):
                # Join list elements with spaces
                return ' '.join(str(item) for item in value if item)
            elif value is not None:
                return str(value)
            return None
        
        # First check if we already extracted a positive prompt (A1111 format)
        if parsed_data.get('positive_prompt'):
            return parsed_data['positive_prompt']
        
        # Check if prompt is already a string
        if isinstance(parsed_data.get('prompt'), str):
            return parsed_data['prompt']
        
        # Check if prompt is a simple value that can be converted
        if parsed_data.get('prompt') and not isinstance(parsed_data.get('prompt'), dict):
            return safe_to_string(parsed_data['prompt'])
        
        prompt_data = parsed_data.get('prompt')
        if isinstance(prompt_data, dict):
            # Use the enhanced logic from parse-metadata.py
            positive_prompt = self._extract_positive_prompt_from_comfyui_data(prompt_data)
            if positive_prompt:
                return positive_prompt
        
        # Check workflow data if available
        workflow_data = parsed_data.get('workflow')
        if isinstance(workflow_data, dict):
            positive_prompt = self._extract_positive_prompt_from_comfyui_data(workflow_data)
            if positive_prompt:
                return positive_prompt
        
        # Check parameters for positive prompt
        if parsed_data.get('parameters', {}).get('positive'):
            return safe_to_string(parsed_data['parameters']['positive'])
        
        return None
    
    def _extract_positive_prompt_from_comfyui_data(self, data):
        """Extract positive prompt from ComfyUI data using the logic from parse-metadata.py."""
        if not isinstance(data, dict):
            return None
        
        # Build nodes dictionary similar to parse-metadata.py
        nodes_by_id = {}
        if "nodes" in data:
            # Handle nodes array format
            for node in data["nodes"]:
                nid = node.get("id")
                if nid is not None:
                    nodes_by_id[nid] = node
        else:
            # Handle flat dictionary format (node_id -> node_data)
            for nid_str, node in data.items():
                try:
                    nid = int(nid_str)
                except:
                    nid = nid_str
                if isinstance(node, dict):
                    if "id" in node:
                        nid = node["id"]
                    nodes_by_id[nid] = node
        
        if not nodes_by_id:
            return None
        
        # First, try to find positive/negative connection pattern
        pos_id = None
        for node in nodes_by_id.values():
            inputs = node.get("inputs", {})
            if "positive" in inputs and "negative" in inputs:
                try:
                    pos_id = int(inputs["positive"][0])
                    break
                except:
                    continue
        
        # Get text from the positive node
        if pos_id is not None and pos_id in nodes_by_id:
            text_val = nodes_by_id[pos_id].get("inputs", {}).get("text")
            if isinstance(text_val, str):
                return text_val
        
        # Fallback: find any text encoder node with text content
        text_encoder_types = [
            'CLIPTextEncode', 'CLIPTextEncodeSDXL', 'CLIPTextEncodeSDXLRefiner',
            'PromptManager', 'BNK_CLIPTextEncoder', 'Text Encoder', 'CLIP Text Encode'
        ]
        
        for node in nodes_by_id.values():
            if isinstance(node, dict):
                class_type = node.get('class_type', '')
                inputs = node.get('inputs', {})
                
                # Check if this is a text encoder node
                if any(encoder_type.lower() in class_type.lower() for encoder_type in text_encoder_types):
                    if 'text' in inputs and isinstance(inputs['text'], str):
                        return inputs['text']
        
        # Final fallback: collect all text fields and return the first non-empty one
        text_fields = []
        for node in nodes_by_id.values():
            if isinstance(node, dict):
                inputs = node.get("inputs", {})
                text_val = inputs.get("text")
                if isinstance(text_val, str) and text_val.strip():
                    text_fields.append(text_val)
        
        # Return the first text field (likely positive prompt)
        if text_fields:
            return text_fields[0]
        
        return None

    # Logging API endpoints
    async def get_logs(self, request):
        """
        Get recent log entries.
        GET /prompt_manager/logs?limit=100&level=INFO
        """
        try:
            # Import logger here to avoid circular imports
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            
            # Get query parameters
            limit = int(request.query.get('limit', 100))
            level = request.query.get('level', None)
            
            # Validate limit
            if limit > 1000:
                limit = 1000
            elif limit < 1:
                limit = 1
            
            # Get logs from memory buffer
            logs = logger_manager.get_recent_logs(limit=limit, level=level)
            
            return web.json_response({
                'success': True,
                'logs': logs,
                'count': len(logs),
                'level_filter': level,
                'limit': limit
            })
            
        except Exception as e:
            self.logger.error(f"Get logs error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e),
                'logs': []
            }, status=500)

    async def get_log_files(self, request):
        """
        Get information about available log files.
        GET /prompt_manager/logs/files
        """
        try:
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            log_files = logger_manager.get_log_files()
            
            return web.json_response({
                'success': True,
                'files': log_files,
                'count': len(log_files)
            })
            
        except Exception as e:
            self.logger.error(f"Get log files error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e),
                'files': []
            }, status=500)

    async def download_log_file(self, request):
        """
        Download a specific log file.
        GET /prompt_manager/logs/download/{filename}
        """
        try:
            filename = request.match_info['filename']
            
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            
            # Validate filename for security
            if not filename or '..' in filename or '/' in filename or '\\' in filename:
                return web.json_response({
                    'success': False,
                    'error': 'Invalid filename'
                }, status=400)
            
            # Get the file content
            log_file_path = logger_manager.log_dir / filename
            
            if not log_file_path.exists():
                return web.json_response({
                    'success': False,
                    'error': 'Log file not found'
                }, status=404)
            
            # Read file content
            with open(log_file_path, 'rb') as f:
                file_content = f.read()
            
            return web.Response(
                body=file_content,
                content_type='text/plain',
                headers={
                    'Content-Disposition': f'attachment; filename="{filename}"',
                    'Content-Length': str(len(file_content))
                }
            )
            
        except Exception as e:
            self.logger.error(f"Download log file error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def truncate_logs(self, request):
        """
        Truncate all log files.
        POST /prompt_manager/logs/truncate
        """
        try:
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            results = logger_manager.truncate_logs()
            
            return web.json_response({
                'success': True,
                'message': f"Truncated {len(results['truncated'])} log files",
                'results': results
            })
            
        except Exception as e:
            self.logger.error(f"Truncate logs error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def get_log_config(self, request):
        """
        Get current logging configuration.
        GET /prompt_manager/logs/config
        """
        try:
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            config = logger_manager.get_config()
            
            return web.json_response({
                'success': True,
                'config': config
            })
            
        except Exception as e:
            self.logger.error(f"Get log config error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def update_log_config(self, request):
        """
        Update logging configuration.
        POST /prompt_manager/logs/config
        Body: {"level": "DEBUG", "console_logging": true, ...}
        """
        try:
            data = await request.json()
            
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            
            # Validate level if provided
            if 'level' in data:
                valid_levels = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']
                if data['level'].upper() not in valid_levels:
                    return web.json_response({
                        'success': False,
                        'error': f'Invalid log level. Must be one of: {valid_levels}'
                    }, status=400)
                data['level'] = data['level'].upper()
            
            logger_manager.update_config(data)
            
            return web.json_response({
                'success': True,
                'message': 'Logging configuration updated',
                'config': logger_manager.get_config()
            })
            
        except Exception as e:
            self.logger.error(f"Update log config error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)

    async def get_log_stats(self, request):
        """
        Get logging statistics.
        GET /prompt_manager/logs/stats
        """
        try:
            try:
                from ..utils.logging_config import get_logger_manager
            except ImportError:
                import sys
                current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                sys.path.insert(0, current_dir)
                from utils.logging_config import get_logger_manager
            
            logger_manager = get_logger_manager()
            stats = logger_manager.get_log_stats()
            
            return web.json_response({
                'success': True,
                'stats': stats
            })
            
        except Exception as e:
            self.logger.error(f"Get log stats error: {e}")
            return web.json_response({
                'success': False,
                'error': str(e)
            }, status=500)