cutThenThink/src/gui/main_window.py

"""
简化的主窗口
极简 GUI，专注核心功能：截图、上传、浏览
"""
import os
from pathlib import Path
from datetime import datetime
from typing import List, Optional

from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QListWidget, QListWidgetItem, QLabel, QStackedWidget,
    QFileDialog, QMessageBox, QStatusBar, QToolBar,
    QStyle, QLineEdit, QComboBox, QDialog, QDialogButtonBox,
    QTextEdit, QCheckBox, QSpinBox, QGroupBox, QFormLayout
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer, QSize
from PyQt6.QtGui import QPixmap, QIcon, QAction, QKeySequence, QShortcut

from src.config import get_config
from src.core.database import Database, Record, get_db
from src.core.screenshot import Screenshot
from src.core.uploader import upload_file, UploadResult
from src.plugins.ocr import get_ocr_plugin


class UploadThread(QThread):
    """上传线程"""
    finished = pyqtSignal(bool, str, str)  # success, url, error

    def __init__(self, filepath: str, provider: str, config: dict):
        super().__init__()
        self.filepath = filepath
        self.provider = provider
        self.config = config

    def run(self):
        """执行上传"""
        try:
            result = upload_file(self.filepath, self.provider, self.config)
            self.finished(result.success, result.url or '', result.error or '')
        except Exception as e:
            self.finished(False, '', str(e))


class OCRThread(QThread):
    """OCR 线程"""
    finished = pyqtSignal(bool, str, str)  # success, text, error

    def __init__(self, filepath: str, ocr_plugin):
        super().__init__()
        self.filepath = filepath
        self.ocr_plugin = ocr_plugin

    def run(self):
        """执行 OCR"""
        success, text, error = self.ocr_plugin.recognize(self.filepath)
        self.finished(success, text or '', error or '')


class SettingsDialog(QDialog):
    """设置对话框"""

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("设置")
        self.setMinimumWidth(500)
        self.config = get_config().app
        self._setup_ui()

    def _setup_ui(self):
        layout = QVBoxLayout(self)

        # 上传设置
        upload_group = QGroupBox("上传设置")
        upload_layout = QFormLayout()

        self.provider_combo = QComboBox()
        self.provider_combo.addItems(["custom", "telegraph", "imgur"])
        self.provider_combo.setCurrentText(self.config.upload.provider)

        self.endpoint_edit = QLineEdit(self.config.upload.endpoint)
        self.endpoint_edit.setPlaceholderText("https://your-server.com/upload")

        self.api_key_edit = QLineEdit(self.config.upload.api_key)
        self.api_key_edit.setPlaceholderText("可选的 API Key")

        self.auto_copy_check = QCheckBox("上传后自动复制链接")
        self.auto_copy_check.setChecked(self.config.upload.auto_copy)

        upload_layout.addRow("上传方式:", self.provider_combo)
        upload_layout.addRow("上传端点:", self.endpoint_edit)
        upload_layout.addRow("API Key:", self.api_key_edit)
        upload_layout.addRow("", self.auto_copy_check)
        upload_group.setLayout(upload_layout)

        # 截图设置
        shot_group = QGroupBox("截图设置")
        shot_layout = QFormLayout()

        self.format_combo = QComboBox()
        self.format_combo.addItems(["png", "jpg", "webp"])
        self.format_combo.setCurrentText(self.config.screenshot.format)

        self.path_edit = QLineEdit(self.config.screenshot.save_path)
        self.path_edit.setPlaceholderText("~/Pictures/Screenshots")

        shot_layout.addRow("保存格式:", self.format_combo)
        shot_layout.addRow("保存路径:", self.path_edit)
        shot_group.setLayout(shot_layout)

        # OCR 设置
        ocr_group = QGroupBox("OCR 设置（可选）")
        ocr_layout = QFormLayout()

        self.ocr_enabled_check = QCheckBox("启用 OCR")
        self.ocr_enabled_check.setChecked(self.config.ocr.enabled)

        self.ocr_auto_copy_check = QCheckBox("识别后自动复制文本")
        self.ocr_auto_copy_check.setChecked(self.config.ocr.auto_copy)

        ocr_layout.addRow("", self.ocr_enabled_check)
        ocr_layout.addRow("", self.ocr_auto_copy_check)
        ocr_group.setLayout(ocr_layout)

        layout.addWidget(upload_group)
        layout.addWidget(shot_group)
        layout.addWidget(ocr_group)

        # 按钮
        buttons = QDialogButtonBox(
            QDialogButtonBox.Ok | QDialogButtonBox.Cancel
        )
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)

    def get_config(self) -> dict:
        """获取配置"""
        return {
            'upload': {
                'provider': self.provider_combo.currentText(),
                'endpoint': self.endpoint_edit.text(),
                'api_key': self.api_key_edit.text(),
                'auto_copy': self.auto_copy_check.isChecked(),
            },
            'screenshot': {
                'format': self.format_combo.currentText(),
                'save_path': self.path_edit.text(),
            },
            'ocr': {
                'enabled': self.ocr_enabled_check.isChecked(),
                'auto_copy': self.ocr_auto_copy_check.isChecked(),
            }
        }


class MainWindow(QMainWindow):
    """主窗口"""

    def __init__(self):
        super().__init__()
        self.config = get_config()
        self.db = get_db()
        self.screenshot = Screenshot(
            self.config.app.screenshot.save_path,
            self.config.app.screenshot.format
        )
        self.ocr_plugin = get_ocr_plugin()
        self.current_records: List[Record] = []
        self.upload_thread: Optional[UploadThread] = None
        self.ocr_thread: Optional[OCRThread] = None
        self._setup_ui()
        self._setup_shortcuts()
        self._load_records()

    def _setup_ui(self):
        """设置 UI"""
        self.setWindowTitle("CutThenThink")
        self.setMinimumSize(900, 600)

        # 创建中心部件
        central = QWidget()
        self.setCentralWidget(central)
        layout = QHBoxLayout(central)

        # 侧边栏
        sidebar = self._create_sidebar()
        layout.addWidget(sidebar, 1)

        # 主内容区
        self.content_stack = QStackedWidget()

        # 记录列表页
        self.records_page = self._create_records_page()
        self.content_stack.addWidget(self.records_page)

        layout.addWidget(self.content_stack, 3)

        # 状态栏
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_bar.showMessage("就绪")

        # 工具栏
        toolbar = QToolBar("主工具栏")
        toolbar.setMovable(False)
        self.addToolBar(toolbar)

        # 截图按钮
        capture_action = QAction("截全屏", self)
        capture_action.setShortcut(QKeySequence(self.config.app.hotkeys.capture))
        capture_action.triggered.connect(self.capture_fullscreen)
        toolbar.addAction(capture_action)

        # 区域截图按钮
        region_action = QAction("区域截图", self)
        region_action.setShortcut(QKeySequence(self.config.app.hotkeys.region))
        region_action.triggered.connect(self.capture_region)
        toolbar.addAction(region_action)

        # 上传按钮
        upload_action = QAction("上传最后截图", self)
        upload_action.setShortcut(QKeySequence(self.config.app.hotkeys.upload))
        upload_action.triggered.connect(self.upload_last)
        toolbar.addAction(upload_action)

        # 设置按钮
        settings_action = QAction("设置", self)
        settings_action.triggered.connect(self.show_settings)
        toolbar.addAction(settings_action)

    def _create_sidebar(self) -> QWidget:
        """创建侧边栏"""
        sidebar = QWidget()
        layout = QVBoxLayout(sidebar)

        # 全部记录
        all_btn = QPushButton("全部记录")
        all_btn.clicked.connect(lambda: self._filter_records(None))
        layout.addWidget(all_btn)

        layout.addStretch()

        return sidebar

    def _create_records_page(self) -> QWidget:
        """创建记录列表页"""
        page = QWidget()
        layout = QVBoxLayout(page)

        # 记录列表
        self.records_list = QListWidget()
        self.records_list.setIconSize(QSize(48, 48))
        self.records_list.itemDoubleClicked.connect(self._open_record_detail)
        layout.addWidget(self.records_list)

        # 操作按钮
        btn_layout = QHBoxLayout()

        refresh_btn = QPushButton("刷新")
        refresh_btn.clicked.connect(self._load_records)
        btn_layout.addWidget(refresh_btn)

        upload_btn = QPushButton("上传选中")
        upload_btn.clicked.connect(self._upload_selected)
        btn_layout.addWidget(upload_btn)

        delete_btn = QPushButton("删除")
        delete_btn.clicked.connect(self._delete_selected)
        btn_layout.addWidget(delete_btn)

        layout.addLayout(btn_layout)

        return page

    def _setup_shortcuts(self):
        """设置快捷键"""
        # ESC 退出
        esc_shortcut = QShortcut(QKeySequence(Qt.Key_Escape), self)
        esc_shortcut.activated.connect(self.close)

    def _load_records(self, category: Optional[str] = None):
        """加载记录"""
        self.current_records = self.db.get_all(category=category)
        self._update_records_list()

    def _update_records_list(self):
        """更新记录列表"""
        self.records_list.clear()

        for record in self.current_records:
            item = QListWidgetItem(record.filename)
            # 设置图标（缩略图）
            if Path(record.filepath).exists():
                pixmap = QPixmap(record.filepath).scaled(
                    48, 48, Qt.KeepAspectRatio, Qt.SmoothTransformation
                )
                item.setIcon(QIcon(pixmap))

            # 添加上传状态标记
            if record.upload_url:
                item.setText(f"✓ {record.filename}")
            else:
                item.setText(f"  {record.filename}")

            item.setData(Qt.UserRole, record)
            self.records_list.addItem(item)

        self.status_bar.showMessage(f"共 {len(self.current_records)} 条记录")

    def _filter_records(self, category: Optional[str]):
        """筛选记录"""
        self._load_records(category)

    def capture_fullscreen(self):
        """全屏截图"""
        filepath, error = self.screenshot.capture_fullscreen()
        self._handle_screenshot_result(filepath, error)

    def capture_region(self):
        """区域截图"""
        pixmap, filepath = self.screenshot.capture_region()
        if pixmap and filepath:
            self._handle_screenshot_result(filepath, None)
        else:
            self.status_bar.showMessage("截图已取消")

    def _handle_screenshot_result(self, filepath: Optional[str], error: Optional[str]):
        """处理截图结果"""
        if error:
            QMessageBox.warning(self, "截图失败", error)
            return

        if not filepath:
            return

        self.status_bar.showMessage(f"截图已保存: {filepath}")

        # 保存到数据库
        file_size = os.path.getsize(filepath)
        record = Record(
            id=0,  # 临时 ID
            filename=Path(filepath).name,
            filepath=filepath,
            file_size=file_size
        )
        record_id = self.db.add(record)

        # 如果启用了 OCR，执行识别
        if self.config.app.ocr.enabled and self.ocr_plugin.is_available():
            self._run_ocr(filepath, record_id)
        else:
            self._load_records()

    def _run_ocr(self, filepath: str, record_id: int):
        """运行 OCR"""
        self.status_bar.showMessage("正在识别文字...")

        self.ocr_thread = OCRThread(filepath, self.ocr_plugin)
        self.ocr_thread.finished.connect(
            lambda success, text, error: self._handle_ocr_result(record_id, success, text, error)
        )
        self.ocr_thread.start()

    def _handle_ocr_result(self, record_id: int, success: bool, text: str, error: str):
        """处理 OCR 结果"""
        if success:
            self.db.update(record_id, ocr_text=text)
            self.status_bar.showMessage("OCR 识别完成")

            if self.config.app.ocr.auto_copy:
                QApplication.clipboard().setText(text)
                self.status_bar.showMessage("已复制 OCR 文本")
        else:
            self.status_bar.showMessage(f"OCR 失败: {error}")

        self._load_records()

    def _upload_selected(self):
        """上传选中的记录"""
        current_item = self.records_list.currentItem()
        if not current_item:
            return

        record = current_item.data(Qt.UserRole)
        self._upload_record(record)

    def _upload_record(self, record: Record):
        """上传单条记录"""
        self.status_bar.showMessage(f"正在上传: {record.filename}")

        self.upload_thread = UploadThread(
            record.filepath,
            self.config.app.upload.provider,
            {
                'endpoint': self.config.app.upload.endpoint,
                'api_key': self.config.app.upload.api_key,
            }
        )
        self.upload_thread.finished.connect(
            lambda success, url, error: self._handle_upload_result(record.id, success, url, error)
        )
        self.upload_thread.start()

    def upload_last(self):
        """上传最后的截图"""
        if not self.current_records:
            QMessageBox.info(self, "提示", "没有可上传的记录")
            return

        # 上传最新的未上传记录
        for record in self.current_records:
            if not record.upload_url:
                self._upload_record(record)
                break

    def _handle_upload_result(self, record_id: int, success: bool, url: str, error: str):
        """处理上传结果"""
        if success:
            self.db.update(record_id, upload_url=url, uploaded_at=datetime.now().isoformat())
            self.status_bar.showMessage(f"上传成功: {url}")

            if self.config.app.upload.auto_copy:
                QApplication.clipboard().setText(url)
                self.status_bar.showMessage("已复制上传链接")
        else:
            QMessageBox.critical(self, "上传失败", error)
            self.status_bar.showMessage("上传失败")

        self._load_records()

    def _delete_selected(self):
        """删除选中的记录"""
        current_item = self.records_list.currentItem()
        if not current_item:
            return

        record = current_item.data(Qt.UserRole)

        reply = QMessageBox.question(
            self, "确认删除",
            f"确定要删除 \"{record.filename}\" 吗？",
            QMessageBox.Yes | QMessageBox.No
        )

        if reply == QMessageBox.Yes:
            self.db.delete(record.id)
            self._load_records()

    def _open_record_detail(self, item: QListWidgetItem):
        """打开记录详情"""
        record = item.data(Qt.UserRole)

        dialog = QDialog(self)
        dialog.setWindowTitle(record.filename)
        dialog.setMinimumSize(500, 400)
        layout = QVBoxLayout(dialog)

        # 图片预览
        if Path(record.filepath).exists():
            pixmap = QPixmap(record.filepath)
            label = QLabel()
            label.setPixmap(pixmap.scaled(
                450, 300, Qt.KeepAspectRatio, Qt.SmoothTransformation
            ))
            layout.addWidget(label)

        # OCR 文本
        if record.ocr_text:
            text_edit = QTextEdit()
            text_edit.setPlainText(record.ocr_text)
            text_edit.setReadOnly(True)
            layout.addWidget(QLabel("识别文字："))
            layout.addWidget(text_edit)

        # 上传链接
        if record.upload_url:
            url_layout = QHBoxLayout()
            url_edit = QLineEdit(record.upload_url)
            url_edit.setReadOnly(True)
            copy_btn = QPushButton("复制")
            copy_btn.clicked.connect(lambda: QApplication.clipboard().setText(record.upload_url))
            url_layout.addWidget(url_edit)
            url_layout.addWidget(copy_btn)
            layout.addWidget(QLabel("上传链接："))
            layout.addLayout(url_layout)

        dialog.exec()

    def show_settings(self):
        """显示设置对话框"""
        dialog = SettingsDialog(self)
        if dialog.exec() == QDialog.Accepted:
            config_data = dialog.get_config()

            # 更新配置
            self.config.app.upload.provider = config_data['upload']['provider']
            self.config.app.upload.endpoint = config_data['upload']['endpoint']
            self.config.app.upload.api_key = config_data['upload']['api_key']
            self.config.app.upload.auto_copy = config_data['upload']['auto_copy']

            self.config.app.screenshot.format = config_data['screenshot']['format']
            self.config.app.screenshot.save_path = config_data['screenshot']['save_path']

            self.config.app.ocr.enabled = config_data['ocr']['enabled']
            self.config.app.ocr.auto_copy = config_data['ocr']['auto_copy']

            # 保存配置
            self.config.save()

            # 更新截图器
            self.screenshot = Screenshot(
                self.config.app.screenshot.save_path,
                self.config.app.screenshot.format
            )

            QMessageBox.information(self, "设置", "设置已保存，部分设置需要重启生效")

    def closeEvent(self, event):
        """关闭事件"""
        self.db.close()
        event.accept()


def main():
    """主函数"""
    import sys
    app = QApplication(sys.argv)
    app.setStyle("Fusion")

    window = MainWindow()
    window.show()

    sys.exit(app.exec())


if __name__ == "__main__":
    main()