ComfyGallery/watcher/fs_watcher.py

# watcher/fs_watcher.py
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 database.db_manager import DBManager
from metadata.parser import MetadataParser

logger = logging.getLogger("ComfyGallery.Watcher")

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
        self._supported_extensions = ('.png', '.jpg', '.jpeg', '.webp')

    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 self._supported_extensions:
            return
        
        # Пауза, гарантирующая полное высвобождение файлового дескриптора при записи из ComfyUI
        time.sleep(0.3)
        try:
            stat = path.stat()
            size = stat.st_size
            mtime = stat.st_mtime
            parent_path = path.parent.resolve()
            folder_id = self._get_or_create_folder_id(parent_path)
            if folder_id is not None:
                file_id = self.db_manager.register_file(
                    folder_id=folder_id, filename=path.name,
                    filepath=str(path.as_posix()), size=size, mtime=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)
                    # Мгновенная реактивная отправка сигнала о добавлении файла в GUI
                    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())
        if not normalized_path.startswith(str(self.root_path.as_posix())):
            return None
        parent_path = folder_path.parent
        parent_id = None
        if parent_path != folder_path and normalized_path != str(self.root_path.as_posix()):
            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 = []