long-screen-cut/gui.py

"""
滚动截屏OCR工具 - PySide6 GUI界面
美观现代的界面设计
"""

import sys
import time
from pathlib import Path
from typing import Optional
from datetime import datetime

from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QTextEdit, QLabel, QFrame, QProgressBar,
    QSystemTrayIcon, QMenu, QStyle
)
from PySide6.QtCore import Qt, QThread, Signal, QTimer, QSize
from PySide6.QtGui import QFont, QIcon, QColor, QPalette, QFontDatabase

from loguru import logger

# 导入核心逻辑
from main import ScrollCaptureOCR, Config


class LogHandler:
    """日志处理器，将日志输出到GUI"""
    def __init__(self, signal):
        self.signal = signal

    def write(self, message):
        if message.strip():
            self.signal.emit(message.strip())

    def flush(self):
        pass


class CaptureWorker(QThread):
    """后台工作线程，执行截屏OCR任务"""
    log_signal = Signal(str)
    progress_signal = Signal(int)
    status_signal = Signal(str)
    finished_signal = Signal()
    error_signal = Signal(str)
    div_processed_signal = Signal(int, int)  # (当前div序号, 总div数)

    def __init__(self, scroll_capture: ScrollCaptureOCR):
        super().__init__()
        self.scroll_capture = scroll_capture
        self.is_running = False

    def run(self):
        """运行截屏任务"""
        self.is_running = True
        try:
            self.status_signal.emit("运行中")

            # 重置状态
            self.scroll_capture.previous_ocr_result = []
            self.scroll_capture.scroll_count = 0
            self.scroll_capture.all_results = []
            self.scroll_capture.processed_divs = []
            self.scroll_capture.last_div_signature = None
            self.scroll_capture.total_scroll_distance = 0
            self.scroll_capture.is_first_capture = True

            # 循环处理
            while self.is_running and self.scroll_capture.process_once():
                progress = min(self.scroll_capture.scroll_count * 10, 90)
                self.progress_signal.emit(progress)
                self.log_signal.emit(f"✓ 第 {self.scroll_capture.scroll_count} 次截屏完成，"
                                   f"累计滚动 {self.scroll_capture.total_scroll_distance} 像素")

            # 保存最终结果
            if self.scroll_capture.all_results:
                self.scroll_capture.save_final_result()
                total_divs = sum(len(result.get('texts', [])) for result in self.scroll_capture.all_results)
                self.log_signal.emit(f"✓ 共处理 {len(self.scroll_capture.all_results)} 次截屏，"
                                   f"识别 {total_divs} 个内容区域")

            self.progress_signal.emit(100)
            self.status_signal.emit("完成")

        except Exception as e:
            self.error_signal.emit(str(e))
            self.status_signal.emit("错误")
        finally:
            self.is_running = False
            self.finished_signal.emit()

    def stop(self):
        """停止任务"""
        self.is_running = False
        self.status_signal.emit("已停止")


class ModernButton(QPushButton):
    """现代风格按钮"""
    def __init__(self, text, parent=None, primary=True):
        super().__init__(text, parent)
        self.primary = primary
        self.setMinimumHeight(40)
        self.setFont(QFont("Microsoft YaHei", 11, QFont.Bold if primary else QFont.Normal))
        self.setCursor(Qt.PointingHandCursor)
        self.update_style()

    def update_style(self):
        if self.primary:
            self.setStyleSheet("""
                QPushButton {
                    background-color: #4CAF50;
                    color: white;
                    border: none;
                    border-radius: 8px;
                    padding: 10px 30px;
                }
                QPushButton:hover {
                    background-color: #45a049;
                }
                QPushButton:pressed {
                    background-color: #3d8b40;
                }
                QPushButton:disabled {
                    background-color: #cccccc;
                    color: #666666;
                }
            """)
        else:
            self.setStyleSheet("""
                QPushButton {
                    background-color: #f5f5f5;
                    color: #333333;
                    border: 2px solid #dddddd;
                    border-radius: 8px;
                    padding: 10px 30px;
                }
                QPushButton:hover {
                    background-color: #e8e8e8;
                    border-color: #cccccc;
                }
                QPushButton:pressed {
                    background-color: #d8d8d8;
                }
            """)


class LogTextEdit(QTextEdit):
    """带样式的日志显示框"""
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setReadOnly(True)
        self.setFont(QFont("Consolas", 10))
        self.setStyleSheet("""
            QTextEdit {
                background-color: #1e1e1e;
                color: #d4d4d4;
                border: 1px solid #3e3e3e;
                border-radius: 8px;
                padding: 10px;
            }
        """)
        self.setPlaceholderText("日志信息将显示在这里...")

    def append_log(self, message: str, level: str = "INFO"):
        """添加带颜色的日志"""
        timestamp = datetime.now().strftime("%H:%M:%S")

        color_map = {
            "INFO": "#4CAF50",
            "WARNING": "#FF9800",
            "ERROR": "#F44336",
            "DEBUG": "#2196F3"
        }
        color = color_map.get(level, "#d4d4d4")

        html = f'<span style="color: #666666;">[{timestamp}]</span> <span style="color: {color};">{message}</span>'
        self.append(html)

        # 自动滚动到底部
        scrollbar = self.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())


class MainWindow(QMainWindow):
    """主窗口"""
    def __init__(self):
        super().__init__()
        self.setWindowTitle("滚动截屏OCR工具")
        self.setMinimumSize(800, 600)
        self.resize(900, 700)

        # 初始化核心逻辑
        self.scroll_capture = ScrollCaptureOCR()
        self.worker: Optional[CaptureWorker] = None

        # 设置应用样式
        self.setup_styles()

        # 创建UI
        self.setup_ui()

        # 设置系统托盘
        self.setup_tray()

        # 重定向日志
        self.setup_logging()

    def setup_styles(self):
        """设置应用样式"""
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f8f9fa;
            }
            QLabel {
                color: #333333;
            }
            QFrame {
                border: none;
            }
        """)

    def setup_ui(self):
        """设置UI界面"""
        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(30, 30, 30, 30)

        # === 标题区域 ===
        title_layout = QHBoxLayout()

        title_label = QLabel("滚动截屏OCR工具")
        title_label.setFont(QFont("Microsoft YaHei", 20, QFont.Bold))
        title_label.setStyleSheet("color: #2c3e50;")
        title_layout.addWidget(title_label)

        title_layout.addStretch()

        # 状态标签
        self.status_label = QLabel("就绪")
        self.status_label.setFont(QFont("Microsoft YaHei", 11))
        self.status_label.setStyleSheet("""
            QLabel {
                color: #4CAF50;
                background-color: #e8f5e9;
                padding: 6px 16px;
                border-radius: 16px;
            }
        """)
        title_layout.addWidget(self.status_label)

        main_layout.addLayout(title_layout)

        # === 信息卡片 ===
        info_card = QFrame()
        info_card.setStyleSheet("""
            QFrame {
                background-color: white;
                border-radius: 12px;
                padding: 15px;
            }
        """)
        info_layout = QHBoxLayout(info_card)
        info_layout.setSpacing(30)

        # 使用说明
        help_text = QLabel(
            "<b>使用步骤：</b><br>"
            "1. 点击「开始」按钮<br>"
            "2. 按住鼠标左键拖动选择区域<br>"
            "3. 程序自动分割div并逐个OCR<br>"
            "4. 智能计算滚动距离，自动翻页<br>"
            "5. 完成后点击「结束」按钮"
        )
        help_text.setFont(QFont("Microsoft YaHei", 10))
        help_text.setStyleSheet("color: #555555; line-height: 1.6;")
        info_layout.addWidget(help_text)

        info_layout.addStretch()

        # 配置信息
        config_text = QLabel(
            f"<b>当前配置：</b><br>"
            f"OCR引擎: {Config.OCR_ENGINE}<br>"
            f"灰度阈值: {Config.GRAY_THRESHOLD}<br>"
            f"输出目录: {Config.OUTPUT_DIR}"
        )
        config_text.setFont(QFont("Microsoft YaHei", 10))
        config_text.setStyleSheet("color: #666666;")
        config_text.setAlignment(Qt.AlignRight)
        info_layout.addWidget(config_text)

        main_layout.addWidget(info_card)

        # === 进度条 ===
        self.progress_bar = QProgressBar()
        self.progress_bar.setMaximumHeight(8)
        self.progress_bar.setTextVisible(False)
        self.progress_bar.setStyleSheet("""
            QProgressBar {
                background-color: #e0e0e0;
                border-radius: 4px;
            }
            QProgressBar::chunk {
                background-color: #4CAF50;
                border-radius: 4px;
            }
        """)
        self.progress_bar.setValue(0)
        main_layout.addWidget(self.progress_bar)

        # === 日志区域 ===
        log_label = QLabel("运行日志")
        log_label.setFont(QFont("Microsoft YaHei", 12, QFont.Bold))
        log_label.setStyleSheet("color: #2c3e50; margin-top: 10px;")
        main_layout.addWidget(log_label)

        self.log_text = LogTextEdit()
        main_layout.addWidget(self.log_text)

        # === 按钮区域 ===
        button_layout = QHBoxLayout()
        button_layout.setSpacing(15)

        button_layout.addStretch()

        # 操作按钮（开始/结束 二合一）
        self.action_btn = ModernButton("开始", primary=True)
        self.action_btn.clicked.connect(self.on_action_button_clicked)
        button_layout.addWidget(self.action_btn)

        # 清空日志按钮
        self.clear_btn = ModernButton("清空日志", primary=False)
        self.clear_btn.clicked.connect(self.clear_logs)
        button_layout.addWidget(self.clear_btn)

        button_layout.addStretch()

        main_layout.addLayout(button_layout)

        # === 底部信息 ===
        footer = QLabel("点击「开始」按钮启动 | 输出目录: ./output/")
        footer.setFont(QFont("Microsoft YaHei", 9))
        footer.setStyleSheet("color: #999999; margin-top: 10px;")
        footer.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(footer)

    def setup_tray(self):
        """设置系统托盘"""
        self.tray_icon = QSystemTrayIcon(self)
        self.tray_icon.setIcon(self.style().standardIcon(QStyle.SP_ComputerIcon))

        # 托盘菜单
        tray_menu = QMenu()
        show_action = tray_menu.addAction("显示窗口")
        show_action.triggered.connect(self.show)

        tray_menu.addSeparator()

        start_action = tray_menu.addAction("开始")
        start_action.triggered.connect(self.start_capture)

        stop_action = tray_menu.addAction("结束")
        stop_action.triggered.connect(self.stop_capture)

        tray_menu.addSeparator()

        quit_action = tray_menu.addAction("退出")
        quit_action.triggered.connect(self.quit_app)

        self.tray_icon.setContextMenu(tray_menu)
        self.tray_icon.show()

    def setup_logging(self):
        """设置日志重定向"""
        # 创建自定义日志处理器
        import logging

        class QtHandler(logging.Handler):
            def __init__(self, callback):
                super().__init__()
                self.callback = callback

            def emit(self, record):
                msg = self.format(record)
                self.callback(msg, record.levelname)

        # 配置loguru输出到GUI
        logger.add(self.log_to_gui, format="{message}")

    def log_to_gui(self, message):
        """将日志输出到GUI"""
        # 在主线程中更新UI
        QTimer.singleShot(0, lambda: self._append_log_safe(message))

    def _append_log_safe(self, message: str):
        """安全地添加日志（在主线程中调用）"""
        level = "INFO"
        if "错误" in message or "失败" in message or "✗" in message:
            level = "ERROR"
        elif "警告" in message:
            level = "WARNING"
        elif "完成" in message or "✓" in message:
            level = "INFO"

        self.log_text.append_log(message, level)

    def on_action_button_clicked(self):
        """操作按钮点击事件"""
        if self.action_btn.text() == "开始":
            self.start_capture()
        else:
            self.stop_capture()

    def start_capture(self):
        """开始截屏"""
        # 检查OCR服务
        if not self.scroll_capture.ocr_engine.check_service():
            self.log_text.append_log("✗ OCR服务未运行", "ERROR")
            if Config.OCR_ENGINE == "umi":
                self.log_text.append_log("请先启动Umi-OCR并开启HTTP服务", "WARNING")
                self.log_text.append_log("设置 → HTTP接口 → 启用HTTP服务", "INFO")
            return

        # 选择区域
        self.log_text.append_log("请在屏幕上拖动选择截图区域...", "INFO")
        self.status_label.setText("选择区域")
        self.status_label.setStyleSheet("""
            QLabel {
                color: #FF9800;
                background-color: #fff3e0;
                padding: 6px 16px;
                border-radius: 16px;
            }
        """)

        try:
            self.scroll_capture.capture_region = self.scroll_capture.region_selector.select_region()
        except Exception as e:
            self.log_text.append_log(f"区域选择失败: {e}", "ERROR")
            self.status_label.setText("就绪")
            self.status_label.setStyleSheet("""
                QLabel {
                    color: #4CAF50;
                    background-color: #e8f5e9;
                    padding: 6px 16px;
                    border-radius: 16px;
                }
            """)
            return

        # 启动工作线程
        self.worker = CaptureWorker(self.scroll_capture)
        self.worker.log_signal.connect(lambda msg: self.log_text.append_log(msg, "INFO"))
        self.worker.progress_signal.connect(self.update_progress)
        self.worker.status_signal.connect(self.update_status)
        self.worker.finished_signal.connect(self.on_finished)
        self.worker.error_signal.connect(self.on_error)

        self.worker.start()

        # 更新UI状态 - 按钮变为"结束"
        self.action_btn.setText("结束")
        self.action_btn.update_style()
        self.progress_bar.setValue(0)

    def stop_capture(self):
        """停止截屏"""
        if self.worker and self.worker.isRunning():
            self.worker.stop()
            self.worker.wait(1000)
            self.log_text.append_log("用户手动结束", "WARNING")

        # 更新UI状态 - 按钮恢复为"开始"
        self.action_btn.setText("开始")
        self.action_btn.update_style()
        self.status_label.setText("就绪")
        self.status_label.setStyleSheet("""
            QLabel {
                color: #4CAF50;
                background-color: #e8f5e9;
                padding: 6px 16px;
                border-radius: 16px;
            }
        """)

    def update_progress(self, value: int):
        """更新进度条"""
        self.progress_bar.setValue(value)

    def update_status(self, status: str):
        """更新状态标签"""
        self.status_label.setText(status)

        if status == "运行中":
            self.status_label.setStyleSheet("""
                QLabel {
                    color: #2196F3;
                    background-color: #e3f2fd;
                    padding: 6px 16px;
                    border-radius: 16px;
                }
            """)
        elif status == "完成":
            self.status_label.setStyleSheet("""
                QLabel {
                    color: #4CAF50;
                    background-color: #e8f5e9;
                    padding: 6px 16px;
                    border-radius: 16px;
                }
            """)
        elif status == "错误":
            self.status_label.setStyleSheet("""
                QLabel {
                    color: #F44336;
                    background-color: #ffebee;
                    padding: 6px 16px;
                    border-radius: 16px;
                }
            """)

    def on_finished(self):
        """任务完成回调"""
        # 按钮恢复为"开始"
        self.action_btn.setText("开始")
        self.action_btn.update_style()
        self.log_text.append_log("✓ 截屏OCR任务已完成", "INFO")

    def on_error(self, error_msg: str):
        """错误回调"""
        self.log_text.append_log(f"✗ 错误: {error_msg}", "ERROR")
        # 按钮恢复为"开始"
        self.action_btn.setText("开始")
        self.action_btn.update_style()

    def clear_logs(self):
        """清空日志"""
        self.log_text.clear()
        self.log_text.append_log("日志已清空", "INFO")

    def quit_app(self):
        """退出应用"""
        if self.worker and self.worker.isRunning():
            self.worker.stop()
            self.worker.wait(2000)
        self.tray_icon.hide()
        QApplication.quit()

    def closeEvent(self, event):
        """关闭事件"""
        # 最小化到托盘而不是关闭
        if self.tray_icon.isVisible():
            self.hide()
            self.tray_icon.showMessage(
                "滚动截屏OCR工具",
                "程序已最小化到系统托盘",
                QSystemTrayIcon.Information,
                2000
            )
            event.ignore()
        else:
            event.accept()


def main():
    """入口函数"""
    app = QApplication(sys.argv)

    # 设置应用信息
    app.setApplicationName("滚动截屏OCR工具")
    app.setApplicationVersion("1.0.0")

    # 设置全局字体
    font = QFont("Microsoft YaHei", 10)
    app.setFont(font)

    # 创建并显示主窗口
    window = MainWindow()
    window.show()

    # 显示启动提示
    window.log_text.append_log("程序已启动，点击「开始」按钮开始", "INFO")
    window.log_text.append_log(f"OCR引擎: {Config.OCR_ENGINE}", "INFO")

    sys.exit(app.exec())


if __name__ == "__main__":
    main()