Comfy-Studio/comfy_studio/ui/main_window.py

"""
Главное окно Comfy-Studio: трёхпанельный интерфейс,
контекстное меню, полноэкранный просмотр, интеграция с ComfyUI.
"""

import os
import sys
import logging
import subprocess
from pathlib import Path

from PyQt6.QtWidgets import (
    QMainWindow, QSplitter, QWidget, QHBoxLayout, QVBoxLayout,
    QStatusBar, QMessageBox, QApplication, QMenu, QInputDialog,
    QLineEdit, QPushButton
)
from PyQt6.QtCore import Qt, QModelIndex, QMimeData, QUrl

from comfy_studio.core.config import DEFAULT_WINDOW_WIDTH, DEFAULT_WINDOW_HEIGHT, APP_NAME
from comfy_studio.core.i18n import tr
from comfy_studio.core.settings import load_settings, save_settings
from comfy_studio.db.manager import DBManager
from comfy_studio.services.watcher import FolderWatcher
from comfy_studio.services.comfy_api import ComfyAPIClient, ComfySendWorker
from comfy_studio.ui.theme import DARK_THEME_QSS
from comfy_studio.ui.widgets.left_panel import LeftPanel
from comfy_studio.ui.widgets.center_grid import CenterGrid, FilePathRole, FileIdRole
from comfy_studio.ui.widgets.right_panel import RightPanel
from comfy_studio.ui.widgets.settings_dialog import SettingsDialog
from comfy_studio.ui.widgets.viewer import FullscreenViewer
from comfy_studio.ui.widgets.translatable_text_edit import set_translate_languages

logger = logging.getLogger(__name__)


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"])
        set_translate_languages(self.settings.get("translate_languages", ["ru", "en"]))
        self.viewer = None
        self.comfy_sender = None

        self.setWindowTitle(APP_NAME)
        self.resize(DEFAULT_WINDOW_WIDTH, DEFAULT_WINDOW_HEIGHT)
        self.setStyleSheet(DARK_THEME_QSS)

        self.init_ui()
        self.bind_events()

        self.left_panel.set_tracked_folders(self.settings["tracked_paths"])
        self.left_panel.update_tags(self.db_manager.get_all_tags())
        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(tr("settings_btn"))
        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([240, 760, 380])
        main_layout.addWidget(self.splitter)

        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage(tr("app_ready"))

    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.left_panel.search_changed.connect(self._on_search)
        self.left_panel.filters_changed.connect(self._reload_current)

        # Центральный грид
        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.fullscreen_requested.connect(self._open_fullscreen)
        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.right_panel.favorite_toggled.connect(self._on_favorite_toggled)
        self.right_panel.tag_added.connect(self._on_tag_added)
        self.right_panel.tag_removed.connect(self._on_tag_removed)
        self.right_panel.json_saved.connect(self._on_json_saved)

        self.settings_btn.clicked.connect(self._open_settings)

        # Реактивное обновление от watcher
        self.watcher.signals.file_added.connect(self._on_file_changed_externally)
        self.watcher.signals.file_removed.connect(self._on_file_changed_externally)

    # ────────────────────────────────
    # Загрузка изображений
    # ────────────────────────────────

    def load_folder_images(self, folder_path: str, search_query: str = ""):
        self.current_root_path = folder_path
        self.center_grid.path_label.setText(f"{tr('folder_label')}: {folder_path}")

        filters = self.left_panel.get_filters()
        files = self.db_manager.get_files_in_folder(
            folder_path, search_query,
            favorites_only=filters["favorites_only"],
            tag=filters["tag"],
        )
        self.center_grid.model.set_files(files)
        self.status_bar.showMessage(f"{tr('files_shown')}: {len(files)}")

    def _reload_current(self):
        if self.current_root_path:
            self.load_folder_images(self.current_root_path,
                                    self.left_panel.search_input.text())

    def _on_subfolder_selected(self, index: QModelIndex):
        folder_path = self.left_panel.folder_model.filePath(index)
        self.load_folder_images(folder_path)

    def _on_search(self, query: str):
        if self.current_root_path:
            self.load_folder_images(self.current_root_path, query)

    def _on_image_selected(self):
        indexes = self.center_grid.list_view.selectedIndexes()
        if not indexes:
            self.right_panel.clear_fields()
            return
        file_id = self.center_grid.model.data(indexes[0], FileIdRole)
        if file_id:
            details = self.db_manager.get_file_details(file_id)
            self.right_panel.display_metadata(details)

    def _on_file_changed_externally(self, filepath: str):
        if not self.current_root_path:
            return
        file_dir = str(Path(filepath).parent.resolve().as_posix())
        current_dir = str(Path(self.current_root_path).resolve().as_posix())
        if file_dir == current_dir:
            self._reload_current()

    # ────────────────────────────────
    # Полноэкранный просмотр
    # ────────────────────────────────

    def _open_fullscreen(self, row: int):
        paths = [f["filepath"] for f in self.center_grid.model.files]
        if not paths:
            return
        self.viewer = FullscreenViewer(paths, row)
        self.viewer.open_fullscreen()

    # ────────────────────────────────
    # Контекстное меню и CRUD
    # ────────────────────────────────

    def _on_context_menu(self, position):
        indexes = self.center_grid.list_view.selectedIndexes()
        if not indexes:
            return

        menu = QMenu(self)

        fullscreen_action = menu.addAction(tr("open_fullscreen"))
        fullscreen_action.setEnabled(len(indexes) == 1)

        favorite_action = menu.addAction(tr("toggle_favorite"))

        show_action = menu.addAction(tr("show_on_disk"))
        show_action.setEnabled(len(indexes) == 1)

        menu.addSeparator()

        rename_action = menu.addAction(tr("rename"))
        rename_action.setEnabled(len(indexes) == 1)

        delete_action = menu.addAction(tr("delete_selected"))

        action = menu.exec(self.center_grid.list_view.mapToGlobal(position))

        if action == fullscreen_action:
            self._open_fullscreen(indexes[0].row())

        elif action == favorite_action:
            for idx in indexes:
                file_id = self.center_grid.model.data(idx, FileIdRole)
                details = self.db_manager.get_file_details(file_id)
                if details:
                    self.db_manager.set_favorite(file_id, not details["favorite"])
            self._reload_current()
            self._on_image_selected()

        elif action == show_action:
            filepath = self.center_grid.model.data(indexes[0], 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, tr("rename_title"), f"{tr('rename_prompt')} {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):
        """Открывает файловый менеджер ОС с выделенным файлом."""
        norm_path = os.path.normpath(filepath)
        if not os.path.exists(norm_path):
            return
        try:
            if sys.platform == "win32":
                subprocess.run(f'explorer /select,"{norm_path}"', shell=True)
            elif sys.platform == "darwin":
                subprocess.run(["open", "-R", norm_path])
            else:
                subprocess.run(["xdg-open", os.path.dirname(norm_path)])
        except Exception as e:
            logger.error(f"Не удалось открыть файловый менеджер: {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)
        try:
            os.rename(str(old_path), str(new_filepath))
            self.status_bar.showMessage(f"{tr('renamed_ok')}: {new_filepath.name}")
            self._reload_current()
        except Exception as e:
            logger.error(f"Ошибка переименования {old_path.name}: {e}")
            QMessageBox.critical(self, tr("error"), f"{tr('rename_failed')}: {e}")

    def _on_copy_file(self, filepath: str):
        if os.path.exists(filepath):
            mime_data = QMimeData()
            mime_data.setUrls([QUrl.fromLocalFile(filepath)])
            QApplication.clipboard().setMimeData(mime_data)
            self.status_bar.showMessage(tr("copied_to_clipboard"))

    def _on_delete_files(self, filepaths: list):
        if not filepaths:
            return
        confirm = QMessageBox.question(
            self, tr("delete_confirm_title"),
            tr("delete_confirm_text", n=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
                except Exception as e:
                    logger.error(f"Ошибка удаления файла {path}: {e}")
            self.status_bar.showMessage(f"{tr('deleted_from_disk')}: {deleted_count}")
            self._reload_current()

    # ────────────────────────────────
    # Метаданные / избранное / теги
    # ────────────────────────────────

    def _on_save_metadata(self, file_id: int, payload: dict):
        success = self.db_manager.update_file_details(
            file_id, payload["positive_prompt"], payload["negative_prompt"], payload["rating"]
        )
        if success:
            self.status_bar.showMessage(tr("params_updated"))
            details = self.db_manager.get_file_details(file_id)
            self.right_panel.display_metadata(details)
            self._reload_current()
        else:
            QMessageBox.critical(self, tr("error"), tr("db_save_error"))

    def _on_json_saved(self, file_id: int, field: str, json_text: str):
        """Сохранение отредактированного workflow/prompt JSON в БД."""
        success = self.db_manager.update_json_field(file_id, field, json_text or None)
        if success:
            self.status_bar.showMessage(tr("json_saved"))
            details = self.db_manager.get_file_details(file_id)
            self.right_panel.display_metadata(details)
            self._reload_current()
        else:
            QMessageBox.critical(self, tr("error"), tr("db_save_error"))

    def _on_favorite_toggled(self, file_id: int, favorite: bool):
        self.db_manager.set_favorite(file_id, favorite)
        self._reload_current()

    def _on_tag_added(self, file_id: int, tag: str):
        self.db_manager.add_tag_to_file(file_id, tag)
        details = self.db_manager.get_file_details(file_id)
        self.right_panel.display_metadata(details)
        self.left_panel.update_tags(self.db_manager.get_all_tags())

    def _on_tag_removed(self, file_id: int, tag: str):
        self.db_manager.remove_tag_from_file(file_id, tag)
        details = self.db_manager.get_file_details(file_id)
        self.right_panel.display_metadata(details)
        self.left_panel.update_tags(self.db_manager.get_all_tags())
        self._reload_current()

    # ────────────────────────────────
    # Настройки
    # ────────────────────────────────

    def _open_settings(self):
        dialog = SettingsDialog(self.settings, self)
        if dialog.exec() == SettingsDialog.DialogCode.Accepted:
            save_settings(self.settings)

            self.api_client = ComfyAPIClient(self.settings["comfyui_url"])
            set_translate_languages(self.settings.get("translate_languages", ["ru", "en"]))
            self.left_panel.set_tracked_folders(self.settings["tracked_paths"])

            self.status_bar.showMessage(tr("watcher_restarting"))
            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(tr("settings_applied"))

    # ────────────────────────────────
    # ComfyUI
    # ────────────────────────────────

    def _on_send_to_comfy(self, prompt_json: str):
        self.status_bar.showMessage(tr("sending_prompt"))
        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(tr("prompt_queued"))
            QMessageBox.information(self, tr("success"), message)
        else:
            self.status_bar.showMessage(tr("prompt_send_failed"))
            QMessageBox.critical(self, tr("network_error_title"), message)