Comfy-Studio/comfy_studio/services/watcher.py

"""
Фоновый мониторинг папок генерации (watchdog) с мостом сигналов в GUI-поток PyQt6.
"""

import os
import time
import logging
from pathlib import Path
from typing import Optional, List

from PyQt6.QtCore import QObject, pyqtSignal
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler

from comfy_studio.core.config import SUPPORTED_EXTENSIONS
from comfy_studio.db.manager import DBManager
from comfy_studio.parsers.metadata_parser import MetadataParser

logger = logging.getLogger(__name__)


class WatcherSignals(QObject):
    """Мост между потоками Watchdog и основным GUI-потоком PyQt6."""
    file_added = pyqtSignal(str)     # filepath
    file_removed = pyqtSignal(str)   # filepath


class GalleryEventHandler(FileSystemEventHandler):
    def __init__(self, db_manager: DBManager, root_path: str, signals: WatcherSignals):
        super().__init__()
        self.db_manager = db_manager
        self.root_path = Path(root_path).resolve()
        self.signals = signals

    def on_created(self, event):
        if event.is_directory:
            self._register_folder_recursive(Path(event.src_path))
        else:
            self._handle_file_creation(Path(event.src_path))

    def on_deleted(self, event):
        if event.is_directory:
            self.db_manager.remove_folder_by_path(event.src_path)
        else:
            self.db_manager.remove_file_by_path(event.src_path)
            self.signals.file_removed.emit(event.src_path)

    def on_moved(self, event):
        src_path = Path(event.src_path)
        dest_path = Path(event.dest_path)
        if event.is_directory:
            self.db_manager.remove_folder_by_path(str(src_path))
            self._register_folder_recursive(dest_path)
        else:
            self.db_manager.remove_file_by_path(str(src_path))
            self._handle_file_creation(dest_path)

    def _handle_file_creation(self, path: Path):
        if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
            return

        # Пауза, гарантирующая полное высвобождение файла при записи из ComfyUI
        time.sleep(0.3)
        try:
            stat = path.stat()
            folder_id = self._get_or_create_folder_id(path.parent.resolve())
            if folder_id is None:
                return
            file_id = self.db_manager.register_file(
                folder_id=folder_id, filename=path.name,
                filepath=str(path.as_posix()), size=stat.st_size, mtime=stat.st_mtime
            )
            if file_id:
                prompt_raw, workflow_raw = MetadataParser.extract_raw_metadata(str(path))
                parsed_params = MetadataParser.parse_comfy_parameters(prompt_raw)
                meta_payload = {
                    "prompt_json": prompt_raw, "workflow_json": workflow_raw,
                    **parsed_params
                }
                self.db_manager.save_metadata(file_id, meta_payload)
                self.signals.file_added.emit(str(path.as_posix()))
        except FileNotFoundError:
            pass
        except Exception as e:
            logger.error(f"Не удалось обработать файл {path}: {e}")

    def _get_or_create_folder_id(self, folder_path: Path) -> Optional[int]:
        normalized_path = str(folder_path.resolve().as_posix())
        root_str = str(self.root_path.as_posix())
        if not normalized_path.startswith(root_str):
            return None
        parent_path = folder_path.parent
        parent_id = None
        if parent_path != folder_path and normalized_path != root_str:
            parent_id = self._get_or_create_folder_id(parent_path)
        return self.db_manager.add_folder(normalized_path, parent_id)

    def _register_folder_recursive(self, folder_path: Path):
        self._get_or_create_folder_id(folder_path)
        try:
            for entry in os.scandir(folder_path):
                entry_path = Path(entry.path)
                if entry.is_dir():
                    self._register_folder_recursive(entry_path)
                elif entry.is_file():
                    self._handle_file_creation(entry_path)
        except Exception as e:
            logger.error(f"Ошибка сканирования папки {folder_path}: {e}")


class FolderWatcher:
    """Управляет мониторингом списка директорий (multi-folder tracking)."""

    def __init__(self, db_manager: DBManager):
        self.db_manager = db_manager
        self.signals = WatcherSignals()
        self.observer: Optional[Observer] = None
        self.handlers: List[GalleryEventHandler] = []

    def start_monitoring(self, root_paths: List[str]):
        if self.observer and self.observer.is_alive():
            self.stop_monitoring()

        self.observer = Observer()
        self.handlers = []

        for p in root_paths:
            path = Path(p).resolve()
            path.mkdir(parents=True, exist_ok=True)

            handler = GalleryEventHandler(self.db_manager, str(path), self.signals)
            logger.info(f"Первоначальная индексация: {path}")
            handler._register_folder_recursive(path)

            self.observer.schedule(handler, str(path), recursive=True)
            self.handlers.append(handler)
            logger.info(f"Отслеживание изменений: {path}")

        self.observer.start()

    def stop_monitoring(self):
        if self.observer:
            self.observer.stop()
            self.observer.join()
            self.observer = None
            self.handlers = []