# ui/main_window.py import os import json import logging import subprocess from pathlib import Path import requests from PyQt6.QtWidgets import ( QMainWindow, QSplitter, QWidget, QHBoxLayout, QVBoxLayout, QStatusBar, QMessageBox, QApplication, QMenu, QInputDialog, QLineEdit, QPushButton ) from PyQt6.QtCore import Qt, QModelIndex, QMimeData, QUrl, QThread, pyqtSignal from ui.styles import DARK_THEME_QSS from ui.widgets.left_panel import LeftPanel from ui.widgets.center_grid import CenterGrid from ui.widgets.right_panel import RightPanel from ui.widgets.settings_dialog import SettingsDialog from database.db_manager import DBManager from watcher.fs_watcher import FolderWatcher from utils.api_client import ComfyAPIClient from utils.settings import load_settings, save_settings FilePathRole = Qt.ItemDataRole.UserRole + 2 logger = logging.getLogger("ComfyGallery.MainWindow") class ComfySendWorker(QThread): finished = pyqtSignal(bool, str) # (success, message) def __init__(self, api_client: ComfyAPIClient, prompt_json_str: str): super().__init__() self.api_client = api_client self.prompt_json_str = prompt_json_str def run(self): logger.info("Запущена фоновая задача отправки промта в ComfyUI.") try: prompt_data = json.loads(self.prompt_json_str) payload = {"prompt": prompt_data} url = f"{self.api_client.base_url}/prompt" logger.debug(f"Отправка POST запроса на {url}") response = requests.post(url, json=payload, timeout=5) if response.status_code == 200: res_data = response.json() prompt_id = res_data.get("prompt_id", "Неизвестно") logger.info(f"Задача успешно принята сервером ComfyUI. ID задачи: {prompt_id}") self.finished.emit( True, f"Промт успешно добавлен в очередь генерации ComfyUI!\n\nID задачи: {prompt_id}" ) else: logger.error(f"Сервер ComfyUI отклонил запрос. Код: {response.status_code}, Ответ: {response.text}") self.finished.emit( False, f"Сервер ComfyUI вернул ошибку {response.status_code}:\n{response.text}" ) except json.JSONDecodeError as je: logger.error(f"Не удалось распарсить prompt_json метаданных: {je}") self.finished.emit(False, f"Ошибка метаданных изображения (неверный формат JSON):\n{je}") except requests.exceptions.RequestException as re: logger.error(f"Сетевое исключение при отправке на {self.api_client.base_url}: {re}") self.finished.emit( False, f"Не удалось подключиться к ComfyUI по адресу {self.api_client.base_url}.\n" f"Убедитесь, что сервер ComfyUI запущен на этом порту.\n\nПодробности:\n{re}" ) except Exception as e: logger.error(f"Непредвиденная ошибка в сетевом воркере: {e}") self.finished.emit(False, f"Произошла непредвиденная ошибка при отправке: {e}") class MainWindow(QMainWindow): def __init__(self, db_manager: DBManager, watcher: FolderWatcher): super().__init__() self.db_manager = db_manager self.watcher = watcher self.settings = load_settings() self.current_root_path = None self.api_client = ComfyAPIClient(self.settings["comfyui_url"]) self.setWindowTitle("ComfyGallery") self.resize(1300, 850) self.setStyleSheet(DARK_THEME_QSS) self.init_ui() self.bind_events() self.left_panel.set_tracked_folders(self.settings["tracked_paths"]) logger.info("Главное окно успешно инициализировано.") def init_ui(self): central_widget = QWidget() self.setCentralWidget(central_widget) main_layout = QVBoxLayout(central_widget) main_layout.setContentsMargins(8, 8, 8, 8) top_bar = QHBoxLayout() self.settings_btn = QPushButton("Настройки подключения и папок") top_bar.addWidget(self.settings_btn) top_bar.addStretch() main_layout.addLayout(top_bar) self.splitter = QSplitter(Qt.Orientation.Horizontal) self.left_panel = LeftPanel() self.center_grid = CenterGrid() self.right_panel = RightPanel() self.splitter.addWidget(self.left_panel) self.splitter.addWidget(self.center_grid) self.splitter.addWidget(self.right_panel) self.splitter.setSizes([230, 720, 350]) main_layout.addWidget(self.splitter) self.status_bar = QStatusBar() self.setStatusBar(self.status_bar) self.status_bar.showMessage("Готов к работе") def bind_events(self): self.left_panel.tree_view.clicked.connect(self._on_subfolder_selected) self.left_panel.root_folder_changed.connect(self.load_folder_images) self.center_grid.list_view.selectionModel().selectionChanged.connect(self._on_image_selected) self.center_grid.list_view.copy_pressed.connect(self._on_copy_file) self.center_grid.list_view.delete_pressed.connect(self._on_delete_files) # Контекстное меню self.center_grid.list_view.customContextMenuRequested.connect(self._on_context_menu) self.right_panel.save_clicked.connect(self._on_save_metadata) self.right_panel.send_to_comfy_clicked.connect(self._on_send_to_comfy) self.left_panel.search_changed.connect(self._on_search) self.settings_btn.clicked.connect(self._open_settings) # Автоматическое реактивное обновление self.watcher.signals.file_added.connect(self._on_file_added_externally) self.watcher.signals.file_removed.connect(self._on_file_removed_externally) def load_folder_images(self, folder_path: str, search_query: str = ""): logger.info(f"Загрузка изображений из папки: {folder_path} (фильтр: '{search_query}')") self.current_root_path = folder_path self.center_grid.path_label.setText(f"Папка: {folder_path}") files = self.db_manager.get_files_in_folder(folder_path, search_query) self.center_grid.model.set_files(files) self.status_bar.showMessage(f"Файлов отображено: {len(files)}") def _on_subfolder_selected(self, index: QModelIndex): folder_path = self.left_panel.folder_model.filePath(index) self.load_folder_images(folder_path) def _on_image_selected(self): indexes = self.center_grid.list_view.selectedIndexes() if not indexes: self.right_panel.clear_fields() return idx = indexes[0] file_id = self.center_grid.model.data(idx, Qt.ItemDataRole.UserRole + 1) if file_id: logger.debug(f"Выделен файл с ID: {file_id}") details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) def _on_search(self, query: str): if self.current_root_path: self.load_folder_images(self.current_root_path, query) # --- РЕАКТИВНЫЕ СОБЫТИЯ WATCHDOG --- def _on_file_added_externally(self, filepath: str): logger.debug(f"Событие Watcher: обнаружен новый файл {filepath}") file_dir = str(Path(filepath).parent.resolve().as_posix()) current_dir = str(Path(self.current_root_path).resolve().as_posix()) if self.current_root_path else "" if file_dir == current_dir: self.load_folder_images(self.current_root_path) def _on_file_removed_externally(self, filepath: str): logger.debug(f"Событие Watcher: файл удален {filepath}") file_dir = str(Path(filepath).parent.resolve().as_posix()) current_dir = str(Path(self.current_root_path).resolve().as_posix()) if self.current_root_path else "" if file_dir == current_dir: self.load_folder_images(self.current_root_path) # --- КОНТЕКСТНОЕ МЕНЮ И CRUD --- def _on_context_menu(self, position): indexes = self.center_grid.list_view.selectedIndexes() if not indexes: return menu = QMenu(self) # Новый пункт контекстного меню show_action = menu.addAction("Показать на диске") show_action.setEnabled(len(indexes) == 1) menu.addSeparator() rename_action = menu.addAction("Переименовать") rename_action.setEnabled(len(indexes) == 1) delete_action = menu.addAction("Удалить выбранные") action = menu.exec(self.center_grid.list_view.mapToGlobal(position)) if action == show_action: idx = indexes[0] filepath = self.center_grid.model.data(idx, FilePathRole) self._show_in_explorer(filepath) elif action == rename_action: idx = indexes[0] filepath = self.center_grid.model.data(idx, FilePathRole) filename = self.center_grid.model.data(idx, Qt.ItemDataRole.DisplayRole) new_name, ok = QInputDialog.getText( self, "Переименование", f"Новое имя файла для {filename}:", QLineEdit.EchoMode.Normal, filename ) if ok and new_name.strip() and new_name != filename: self._rename_file(filepath, new_name.strip()) elif action == delete_action: paths = [self.center_grid.model.data(idx, FilePathRole) for idx in indexes] self._on_delete_files(paths) def _show_in_explorer(self, filepath: str): """ Открывает Проводник Windows и подсвечивает (выделяет) данный файл. """ norm_path = os.path.normpath(filepath) if not os.path.exists(norm_path): logger.warning(f"Попытка показать несуществующий файл: {norm_path}") return logger.info(f"Открытие проводника для файла: {norm_path}") try: # Команда explorer /select открывает папку и фокусом выделяет файл subprocess.run(f'explorer /select,"{norm_path}"', shell=True) except Exception as e: logger.error(f"Не удалось открыть Проводник Windows: {e}") def _rename_file(self, old_filepath: str, new_filename: str): old_path = Path(old_filepath) new_filepath = old_path.parent / new_filename if not new_filepath.suffix: new_filepath = new_filepath.with_suffix(old_path.suffix) logger.info(f"Переименование файла: {old_path.name} -> {new_filepath.name}") try: os.rename(str(old_path), str(new_filepath)) self.status_bar.showMessage(f"Файл успешно переименован: {new_filepath.name}") if self.current_root_path: self.load_folder_images(self.current_root_path) except Exception as e: logger.error(f"Ошибка переименования файла {old_path.name}: {e}") QMessageBox.critical(self, "Ошибка", f"Не удалось переименовать файл: {e}") def _on_copy_file(self, filepath: str): logger.info(f"Копирование файла в буфер обмена: {filepath}") if os.path.exists(filepath): mime_data = QMimeData() mime_data.setUrls([QUrl.fromLocalFile(filepath)]) QApplication.clipboard().setMimeData(mime_data) self.status_bar.showMessage("Изображение скопировано в буфер обмена.") def _on_delete_files(self, filepaths: list): if not filepaths: return logger.info(f"Запрос удаления группы файлов: {len(filepaths)} шт.") confirm = QMessageBox.question( self, "Удаление", f"Вы действительно хотите безвозвратно удалить {len(filepaths)} файлов?", QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No ) if confirm == QMessageBox.StandardButton.Yes: deleted_count = 0 for path in filepaths: try: if os.path.exists(path): os.remove(path) deleted_count += 1 logger.debug(f"Файл успешно удален физически: {path}") except Exception as e: logger.error(f"Ошибка физического удаления файла {path}: {e}") self.status_bar.showMessage(f"Удалено файлов с диска: {deleted_count}") if self.current_root_path: self.load_folder_images(self.current_root_path) # --- НАСТРОЙКИ --- def _open_settings(self): logger.info("Открытие диалогового окна настроек.") dialog = SettingsDialog(self.settings, self) if dialog.exec() == SettingsDialog.DialogCode.Accepted: save_settings(self.settings) logger.info("Пользователь сохранил новые настройки.") self.api_client = ComfyAPIClient(self.settings["comfyui_url"]) self.left_panel.set_tracked_folders(self.settings["tracked_paths"]) self.status_bar.showMessage("Перезапуск службы отслеживания...") self.watcher.stop_monitoring() self.watcher.start_monitoring(self.settings["tracked_paths"]) if self.settings["tracked_paths"]: first_path = self.settings["tracked_paths"][0] self.left_panel.set_root_path(first_path) self.load_folder_images(first_path) self.status_bar.showMessage("Настройки успешно применены.") # --- ВЗАИМОДЕЙСТВИЕ С COMFYUI --- def _on_save_metadata(self, file_id: int, payload: dict): logger.info(f"Сохранение измененных пользователем метаданных для ID {file_id}") success = self.db_manager.update_file_details( file_id, payload["positive_prompt"], payload["negative_prompt"], payload["rating"] ) if success: self.status_bar.showMessage("Параметры успешно обновлены.") details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) if self.current_root_path: self.load_folder_images(self.current_root_path) else: logger.error(f"Не удалось обновить метаданные в БД для ID {file_id}") QMessageBox.critical(self, "Ошибка", "Не удалось сохранить изменения в БД.") def _on_send_to_comfy(self, prompt_json: str): self.status_bar.showMessage("Отправка промта в ComfyUI...") self.right_panel.send_btn.setEnabled(False) # Создаем и запускаем фоновый поток воркера self.comfy_sender = ComfySendWorker(self.api_client, prompt_json) self.comfy_sender.finished.connect(self._on_send_finished) self.comfy_sender.start() def _on_send_finished(self, success: bool, message: str): self.right_panel.send_btn.setEnabled(True) if success: self.status_bar.showMessage("Промт успешно добавлен в очередь ComfyUI!") QMessageBox.information(self, "Успешно отправлено", message) else: self.status_bar.showMessage("Не удалось отправить промт.") QMessageBox.critical(self, "Ошибка сети / API", message)