ComfyAssets-ComfyUI_PromptManager/utils/diagnostics.py

"""Diagnostic utilities for troubleshooting the PromptManager gallery system.

This module provides comprehensive diagnostic tools to help identify and resolve
issues with the PromptManager system, particularly the gallery and image monitoring
functionality. It performs systematic checks of all system components and provides
clear feedback on potential problems.

The diagnostic system checks:
- Database connectivity and schema integrity
- Image table structure and data
- File system permissions and access
- ComfyUI output directory detection
- Required Python dependencies
- System integration points

Typical usage:
    from utils.diagnostics import GalleryDiagnostics, run_diagnostics

    # Run full diagnostic suite
    results = run_diagnostics()

    # Or create custom diagnostic instance
    diagnostics = GalleryDiagnostics("custom_db.db")
    database_status = diagnostics.check_database()

The diagnostics provide:
- Clear pass/fail status for each component
- Detailed error messages and remediation suggestions
- Statistics and system information
- Test utilities for verifying functionality
"""

import os
import sqlite3
from contextlib import closing
from pathlib import Path
from typing import Dict, Any

from .logging_config import get_logger

try:
    from ..database.workflow_documents import DocumentReader, resolve_image_row
except ImportError:  # loaded as a top-level ``utils`` package
    from database.workflow_documents import DocumentReader, resolve_image_row


class GalleryDiagnostics:
    """Diagnostics for the gallery system.

    This class provides comprehensive diagnostic capabilities for the PromptManager
    gallery system. It systematically checks all components and dependencies to
    identify potential issues and provide actionable feedback.

    The diagnostics cover:
    - Database structure and connectivity
    - Image tracking table integrity
    - File system access and permissions
    - ComfyUI integration points
    - Python dependency availability

    Each diagnostic method returns a standardized result dictionary with:
    - status: 'ok', 'warning', or 'error'
    - message: Descriptive message (for warnings/errors)
    - Additional data specific to the diagnostic
    """

    def __init__(self, db_path: str = None):
        """Initialize the diagnostics system.

        Args:
            db_path: Path to the SQLite database file to diagnose.
                     If None, resolves from config.json or defaults to
                     prompts.db in the extension root.
        """
        if db_path is None:
            db_path = self._resolve_db_path()
        self.db_path = db_path
        self.logger = get_logger("prompt_manager.diagnostics")

    @staticmethod
    def _resolve_db_path() -> str:
        """Resolve database path from config, matching PromptDatabase logic."""
        extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        try:
            from py.config import PromptManagerConfig

            db_path = PromptManagerConfig.DEFAULT_DB_PATH
        except Exception:
            db_path = "prompts.db"

        if not os.path.isabs(db_path):
            db_path = os.path.normpath(os.path.join(extension_root, db_path))
        return db_path

    def run_full_diagnostic(self) -> Dict[str, Any]:
        """Run a complete diagnostic check.

        Executes all diagnostic checks in sequence and provides a comprehensive
        report of system status. Logs detailed information during the process.

        Returns:
            Dictionary mapping diagnostic categories to their results:
            - database: Database connectivity and structure check
            - images_table: Image tracking table specific check
            - file_system: File system access and permissions check
            - comfyui_output: ComfyUI output directory detection
            - dependencies: Python dependency availability check
        """
        self.logger.info("\n" + "=" * 60)
        self.logger.info("[DIAG] PROMPTMANAGER GALLERY DIAGNOSTICS")
        self.logger.info("=" * 60)

        results = {
            "database": self.check_database(),
            "images_table": self.check_images_table(),
            "file_system": self.check_file_system(),
            "comfyui_output": self.check_comfyui_output(),
            "dependencies": self.check_dependencies(),
        }

        self.logger.info("\n" + "=" * 60)
        self.logger.info("[SUMMARY] DIAGNOSTIC SUMMARY")
        self.logger.info("=" * 60)

        for category, result in results.items():
            status = "[PASS] PASS" if result["status"] == "ok" else "[FAIL] FAIL"
            self.logger.info(f"{category.upper():<20} {status}")
            if result["status"] != "ok":
                self.logger.warning(f"  Issue: {result['message']}")

        self.logger.info("\n" + "=" * 60)
        return results

    def check_database(self) -> Dict[str, Any]:
        """Check database connection and structure.

        Verifies that the database file exists, is accessible, and contains
        the expected prompt table structure.

        Returns:
            Dictionary containing:
            - status: 'ok' or 'error'
            - message: Error description (if status is 'error')
            - prompt_count: Number of prompts in database (if successful)
            - has_images_table: Whether generated_images table exists (if successful)
        """
        self.logger.info("\n[DB]  Checking Database...")

        try:
            if not os.path.exists(self.db_path):
                return {
                    "status": "error",
                    "message": f"Database file not found: {self.db_path}",
                }

            with closing(sqlite3.connect(self.db_path)) as conn:
                conn.row_factory = sqlite3.Row

                # Check prompts table
                cursor = conn.execute("SELECT COUNT(*) as count FROM prompts")
                prompt_count = cursor.fetchone()["count"]
                self.logger.info(f"   [NOTE] Prompts in database: {prompt_count}")

                # Check if generated_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
                self.logger.info(f"   [IMG]  Images table exists: {has_images_table}")

                return {
                    "status": "ok",
                    "prompt_count": prompt_count,
                    "has_images_table": has_images_table,
                }

        except Exception as e:
            return {"status": "error", "message": f"Database error: {str(e)}"}

    def check_images_table(self) -> Dict[str, Any]:
        """Check the generated_images table specifically.

        Examines the generated_images table structure and content to verify
        the gallery system can function properly.

        Returns:
            Dictionary containing:
            - status: 'ok' or 'error'
            - message: Error description (if table missing or inaccessible)
            - image_count: Number of images in table (if successful)
            - recent_images: List of recent image records (if successful)
        """
        self.logger.info("\n[IMG]  Checking Images Table...")

        try:
            with closing(sqlite3.connect(self.db_path)) as conn:
                conn.row_factory = sqlite3.Row

                # Check if table exists
                cursor = conn.execute("""
                    SELECT name FROM sqlite_master
                    WHERE type='table' AND name='generated_images'
                """)

                if not cursor.fetchone():
                    return {
                        "status": "error",
                        "message": (
                            "generated_images table does not exist - "
                            "run the updated code to create it"
                        ),
                    }

                # Check image records
                cursor = conn.execute("SELECT COUNT(*) as count FROM generated_images")
                image_count = cursor.fetchone()["count"]
                self.logger.info(f"   [STATS] Images in database: {image_count}")

                # Get recent images
                cursor = conn.execute("""
                    SELECT gi.*, p.text
                    FROM generated_images gi
                    LEFT JOIN prompts p ON gi.prompt_id = p.id
                    ORDER BY gi.generation_time DESC
                    LIMIT 5
                """)
                reader = DocumentReader(conn)
                recent_images = [resolve_image_row(reader, row) for row in cursor]

                self.logger.info(f"   [TIME] Recent images: {len(recent_images)}")
                for img in recent_images:
                    self.logger.info(
                        f"      - {img['filename']} -> Prompt {img['prompt_id']}"
                    )

                return {
                    "status": "ok",
                    "image_count": image_count,
                    "recent_images": recent_images,
                }

        except Exception as e:
            return {"status": "error", "message": f"Images table error: {str(e)}"}

    def check_file_system(self) -> Dict[str, Any]:
        """Check file system and permissions.

        Verifies that the application has appropriate file system access
        for reading images and writing database files.

        Returns:
            Dictionary containing:
            - status: 'ok' or 'error'
            - message: Error description (if issues detected)
            - current_dir: Current working directory path
            - can_write: Whether write access is available
        """
        self.logger.info("\n[DIR] Checking File System...")

        try:
            # Check current directory
            current_dir = os.getcwd()
            self.logger.info(f"   [FOLDER] Current directory: {current_dir}")

            # Check if we can write to current directory
            test_file = "test_write.tmp"
            try:
                with open(test_file, "w") as f:
                    f.write("test")
                os.remove(test_file)
                can_write = True
            except (IOError, OSError):
                can_write = False

            self.logger.info(f"   [EDIT]  Can write to directory: {can_write}")

            return {"status": "ok", "current_dir": current_dir, "can_write": can_write}

        except Exception as e:
            return {"status": "error", "message": f"File system error: {str(e)}"}

    def check_comfyui_output(self) -> Dict[str, Any]:
        """Check ComfyUI output directories.

        Attempts to locate ComfyUI output directories where generated images
        would be stored. Checks both common locations and ComfyUI's configured paths.

        Returns:
            Dictionary containing:
            - status: 'ok' if directories found, 'warning' if none found
            - message: Warning message (if no directories found)
            - output_dirs: List of detected output directory paths
        """
        self.logger.info("\n[STYLE] Checking ComfyUI Output...")

        output_dirs = []

        def _add_dir(abs_path, source):
            if abs_path not in output_dirs:
                output_dirs.append(abs_path)
                self.logger.info(f"   [DIR] {source}: {abs_path}")
                try:
                    image_files = []
                    for ext in [".png", ".jpg", ".jpeg", ".webp"]:
                        image_files.extend(Path(abs_path).rglob(f"*{ext}"))
                    self.logger.info(f"      [IMG]  Images found: {len(image_files)}")
                except Exception as e:
                    self.logger.error(f"      [FAIL] Error scanning: {e}")

        # Check user-configured gallery directories from config.json
        try:
            from py.config import GalleryConfig

            for cfg_dir in GalleryConfig.MONITORING_DIRECTORIES:
                abs_path = os.path.abspath(cfg_dir)
                if os.path.exists(abs_path):
                    _add_dir(abs_path, "Configured directory")
        except Exception as e:
            self.logger.debug(f"   [NOTE] Could not load gallery config: {e}")

        # Check ComfyUI's own output directory (respects --output-directory)
        try:
            import folder_paths

            comfyui_output = folder_paths.get_output_directory()
            if comfyui_output and os.path.exists(comfyui_output):
                _add_dir(os.path.abspath(comfyui_output), "ComfyUI output dir")
        except ImportError:
            self.logger.warning("   [WARN]  ComfyUI folder_paths not available")

        # Fallback: check common relative paths
        if not output_dirs:
            for dir_path in [
                "output",
                "../output",
                "../../output",
                "ComfyUI/output",
                "../ComfyUI/output",
                "../../ComfyUI/output",
            ]:
                abs_path = os.path.abspath(dir_path)
                if os.path.exists(abs_path):
                    _add_dir(abs_path, "Found output dir")

        return {
            "status": "ok" if output_dirs else "warning",
            "message": "No output directories found" if not output_dirs else None,
            "output_dirs": output_dirs,
        }

    def check_dependencies(self) -> Dict[str, Any]:
        """Check required dependencies.

        Verifies that all required Python packages are available for the
        PromptManager system to function properly.

        Returns:
            Dictionary containing:
            - status: 'ok' if all dependencies available, 'error' if any missing
            - message: Error description (if dependencies missing)
            - dependencies: Dictionary mapping package names to availability status
        """
        self.logger.info("\n[PKG] Checking Dependencies...")

        dependencies = {"watchdog": False, "PIL": False, "sqlite3": False}

        # Check watchdog
        try:
            import watchdog

            dependencies["watchdog"] = True
            version = getattr(watchdog, "__version__", "unknown version")
            self.logger.info(f"   [PASS] watchdog: {version}")
        except ImportError:
            self.logger.error("   [FAIL] watchdog: NOT INSTALLED")

        # Check PIL
        try:
            from PIL import Image  # noqa: F401 - availability probe

            dependencies["PIL"] = True
            self.logger.info("   [PASS] PIL (Pillow): Available")
        except ImportError:
            self.logger.error("   [FAIL] PIL (Pillow): NOT AVAILABLE")

        # Check sqlite3
        try:
            import sqlite3

            dependencies["sqlite3"] = True
            self.logger.info(f"   [PASS] sqlite3: {sqlite3.sqlite_version}")
        except ImportError:
            self.logger.error("   [FAIL] sqlite3: NOT AVAILABLE")

        all_deps_ok = all(dependencies.values())

        return {
            "status": "ok" if all_deps_ok else "error",
            "message": "Missing dependencies" if not all_deps_ok else None,
            "dependencies": dependencies,
        }

    def create_test_image_link(
        self, prompt_id: int, test_image_path: str = None
    ) -> Dict[str, Any]:
        """Create a test image link to verify the system works.

        Creates a test entry in the generated_images table to verify that
        the image linking functionality is working correctly.

        Args:
            prompt_id: ID of an existing prompt to link the test image to
            test_image_path: Optional path for the test image (uses fake path if None)

        Returns:
            Dictionary containing:
            - status: 'ok' if test link created successfully, 'error' otherwise
            - message: Success or error message
            - image_id: ID of created test image record (if successful)
        """
        self.logger.info(f"\n[TEST] Creating test image link for prompt {prompt_id}...")

        try:
            # Import database operations
            import sys
            import os

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

            db = PromptDatabase()

            # Create a fake image path if none provided
            if not test_image_path:
                test_image_path = "/fake/test/image.png"

            # Create test metadata
            test_metadata = {
                "file_info": {
                    "size": 1024000,
                    "dimensions": [512, 512],
                    "format": "PNG",
                },
                "workflow": {"test": True},
                "prompt": {"test_prompt": "This is a test"},
            }

            # Link the test image
            image_id = db.link_image_to_prompt(
                prompt_id=str(prompt_id),
                image_path=test_image_path,
                metadata=test_metadata,
            )

            self.logger.info(f"   [PASS] Test image linked with ID: {image_id}")

            return {
                "status": "ok",
                "image_id": image_id,
                "message": f"Test image linked successfully with ID {image_id}",
            }

        except Exception as e:
            return {
                "status": "error",
                "message": f"Failed to create test link: {str(e)}",
            }


def run_diagnostics() -> Dict[str, Any]:
    """Run diagnostics from command line.

    Convenience function to create a GalleryDiagnostics instance and run
    the full diagnostic suite with default settings.

    Returns:
        Complete diagnostic results dictionary
    """
    diagnostics = GalleryDiagnostics()
    return diagnostics.run_full_diagnostic()


if __name__ == "__main__":
    run_diagnostics()