desktop-transfer/ui/main_window.py

from PySide6.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QTextEdit, QPushButton,
    QLineEdit, QLabel, QGroupBox, QScrollArea, QToolButton, QStatusBar,
    QFileDialog, QMessageBox, QListWidget, QListWidgetItem, QFrame, QSplitter
)
from PySide6.QtCore import Qt, QTimer, QSize
from PySide6.QtGui import QFont, QPalette, QColor
from utils.word_handler import WordHandler
from utils.system_monitor import SystemMonitor
from translator import Translator
from utils.logger import logger
import os

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("PrivaTrans")
        self.setMinimumSize(800, 600)
        
        # 初始化组件
        self.translator = Translator()
        self.system_monitor = SystemMonitor()
        self.word_handler = WordHandler()
        
        # 默认模型路径
        self.default_model_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models", "HY-MT1.5-1.8B_bf16_Q4_K_M.gguf")
        
        # 术语列表
        self.terms = []
        
        # 初始化UI
        self.init_ui()
        
        # 初始化状态栏定时器
        self.init_status_bar_timer()
        
        # 加载默认模型
        self.load_default_model()
    
    def init_ui(self):
        """初始化UI界面"""
        # 主布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.setSpacing(20)
        main_layout.setContentsMargins(20, 20, 20, 10)
        
        # 顶部标题
        title_label = QLabel("PrivaTrans")
        title_font = QFont()
        title_font.setPointSize(24)
        title_font.setBold(True)
        title_label.setFont(title_font)
        title_label.setStyleSheet("color: #1a365d;")
        
        # 模型信息和更换按钮
        model_layout = QHBoxLayout()
        
        self.model_status_label = QLabel()
        self.model_status_label.setStyleSheet("color: #2d3748;")
        
        self.model_info_label = QLabel()
        self.model_info_label.setStyleSheet("color: #2d3748; font-weight: bold;")
        
        self.change_model_btn = QPushButton("更换")
        self.change_model_btn.setFixedSize(80, 32)
        self.change_model_btn.setStyleSheet(
            "background-color: #e2e8f0; color: #2d3748; border: none; border-radius: 6px;"
        )
        self.change_model_btn.clicked.connect(self.change_model)
        
        model_layout.addWidget(self.model_status_label)
        model_layout.addWidget(self.model_info_label)
        model_layout.addStretch()
        model_layout.addWidget(self.change_model_btn)
        
        # 高级辅助面板（可折叠）
        advanced_panel = QWidget()
        advanced_layout = QVBoxLayout(advanced_panel)
        advanced_layout.setSpacing(10)
        advanced_layout.setContentsMargins(0, 0, 0, 0)
        
        # 高级辅助标题和折叠按钮
        advanced_header = QWidget()
        header_layout = QHBoxLayout(advanced_header)
        header_layout.setContentsMargins(0, 0, 0, 0)
        header_layout.setSpacing(10)
        
        advanced_title = QLabel("高级辅助 (背景与术语)")
        advanced_title.setStyleSheet("color: #4a5568; font-weight: bold;")
        
        self.advanced_toggle = QToolButton()
        self.advanced_toggle.setText("▼")
        self.advanced_toggle.setCheckable(True)
        self.advanced_toggle.setChecked(False)
        self.advanced_toggle.setFixedSize(20, 20)
        self.advanced_toggle.setStyleSheet(
            "border: none; background: none; color: #4a5568;"
        )
        self.advanced_toggle.clicked.connect(self.toggle_advanced_panel)
        
        header_layout.addWidget(advanced_title)
        header_layout.addStretch()
        header_layout.addWidget(self.advanced_toggle)
        
        advanced_layout.addWidget(advanced_header)
        
        # 高级辅助内容区域
        self.advanced_content = QWidget()
        content_layout = QVBoxLayout(self.advanced_content)
        content_layout.setSpacing(15)
        content_layout.setContentsMargins(0, 0, 0, 0)
        
        # 文本背景/场景介绍
        context_label = QLabel("文本背景 / 场景介绍")
        context_label.setStyleSheet("color: #4a5568;")
        
        self.context_edit = QTextEdit()
        self.context_edit.setPlaceholderText("例如：这是一份关于建筑工程的合同...")
        self.context_edit.setFixedHeight(80)
        self.context_edit.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 6px; padding: 8px;"
        )
        
        # 术语定义
        terms_label = QLabel("术语簿 (定义 A=B)")
        terms_label.setStyleSheet("color: #4a5568;")
        
        self.terms_list = QListWidget()
        self.terms_list.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 6px;"
        )
        
        terms_input_layout = QHBoxLayout()
        self.term_input = QLineEdit()
        self.term_input.setPlaceholderText("你好 = what's up")
        self.term_input.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 6px; padding: 8px;"
        )
        
        add_term_btn = QPushButton("+")
        add_term_btn.setFixedSize(32, 32)
        add_term_btn.setStyleSheet(
            "background-color: #3182ce; color: white; border: none; border-radius: 6px;"
        )
        add_term_btn.clicked.connect(self.add_term)
        
        terms_input_layout.addWidget(self.term_input)
        terms_input_layout.addWidget(add_term_btn)
        
        # 术语列表项右键菜单
        self.terms_list.setContextMenuPolicy(Qt.CustomContextMenu)
        self.terms_list.customContextMenuRequested.connect(self.show_term_context_menu)
        
        content_layout.addWidget(context_label)
        content_layout.addWidget(self.context_edit)
        content_layout.addWidget(terms_label)
        content_layout.addWidget(self.terms_list)
        content_layout.addLayout(terms_input_layout)
        
        advanced_layout.addWidget(self.advanced_content)
        
        # 默认隐藏高级辅助内容
        self.advanced_content.setVisible(False)  
        
        # 将高级辅助面板添加到主布局
        advanced_panel.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 8px; padding: 15px; margin-top: 10px;"
        )
        
        # 原文输入区域
        input_layout = QVBoxLayout()
        
        input_header_layout = QHBoxLayout()
        input_label = QLabel("原文内容")
        input_label.setStyleSheet("color: #4a5568; font-weight: bold;")
        
        import_word_btn = QPushButton("导入 Word")
        import_word_btn.setFixedSize(120, 32)
        import_word_btn.setStyleSheet(
            "background-color: #e2e8f0; color: #2d3748; border: none; border-radius: 6px;"
        )
        import_word_btn.clicked.connect(self.import_word)
        
        input_header_layout.addWidget(input_label)
        input_header_layout.addStretch()
        input_header_layout.addWidget(import_word_btn)
        
        self.source_text = QTextEdit()
        self.source_text.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 6px; padding: 10px;"
        )
        
        input_layout.addLayout(input_header_layout)
        input_layout.addWidget(self.source_text)
        
        # 翻译按钮
        self.translate_btn = QPushButton("开始翻译")
        self.translate_btn.setFixedHeight(50)
        self.translate_btn.setStyleSheet(
            "background-color: #3182ce; color: white; border: none; border-radius: 8px; font-size: 16px; font-weight: bold;"
        )
        self.translate_btn.clicked.connect(self.start_translation)
        
        # 译文结果区域
        output_layout = QVBoxLayout()
        
        output_header_layout = QHBoxLayout()
        output_label = QLabel("译文结果")
        output_label.setStyleSheet("color: #4a5568; font-weight: bold;")
        
        output_buttons_layout = QHBoxLayout()
        export_word_btn = QPushButton("导出 Word")
        export_word_btn.setFixedSize(120, 32)
        export_word_btn.setStyleSheet(
            "background-color: #e2e8f0; color: #2d3748; border: none; border-radius: 6px; margin-right: 10px;"
        )
        export_word_btn.clicked.connect(self.export_word)
        
        copy_btn = QPushButton("复制内容")
        copy_btn.setFixedSize(120, 32)
        copy_btn.setStyleSheet(
            "background-color: #e2e8f0; color: #2d3748; border: none; border-radius: 6px;"
        )
        copy_btn.clicked.connect(self.copy_result)
        
        output_buttons_layout.addWidget(export_word_btn)
        output_buttons_layout.addWidget(copy_btn)
        
        output_header_layout.addWidget(output_label)
        output_header_layout.addStretch()
        output_header_layout.addLayout(output_buttons_layout)
        
        self.result_text = QTextEdit()
        self.result_text.setReadOnly(True)
        self.result_text.setStyleSheet(
            "border: 1px solid #e2e8f0; border-radius: 6px; padding: 10px;"
        )
        
        output_layout.addLayout(output_header_layout)
        output_layout.addWidget(self.result_text)
        
        # 将所有组件添加到主布局
        main_layout.addWidget(title_label)
        main_layout.addLayout(model_layout)
        main_layout.addWidget(advanced_panel)
        main_layout.addLayout(input_layout)
        main_layout.addWidget(self.translate_btn)
        main_layout.addLayout(output_layout)
        
        # 状态栏
        self.status_bar = QStatusBar()
        self.setStatusBar(self.status_bar)
        self.status_label = QLabel()
        self.status_bar.addWidget(self.status_label)
    
    def init_status_bar_timer(self):
        """初始化状态栏定时器"""
        self.status_timer = QTimer()
        self.status_timer.setInterval(1000)  # 每秒更新一次
        self.status_timer.timeout.connect(self.update_status_bar)
        self.status_timer.start()
    
    def update_status_bar(self):
        """更新状态栏显示"""
        status_text = self.system_monitor.get_status_text()
        self.status_label.setText(status_text)
    
    def load_default_model(self):
        """加载默认模型"""
        if os.path.exists(self.default_model_path):
            self.model_status_label.setText("模型加载中...")
            if self.translator.load_model(self.default_model_path):
                self.model_status_label.setText("● 模型就绪")
                self.model_status_label.setStyleSheet("color: #38a169;")
                self.model_info_label.setText(f"{self.translator.get_model_info()}")
            else:
                self.model_status_label.setText("● 模型加载失败")
                self.model_status_label.setStyleSheet("color: #e53e3e;")
        else:
            self.model_status_label.setText("● 模型文件不存在")
            self.model_status_label.setStyleSheet("color: #e53e3e;")
    
    def change_model(self):
        """更换模型"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择模型文件", "", "GGUF模型文件 (*.gguf)"
        )
        if file_path:
            self.model_status_label.setText("模型加载中...")
            if self.translator.load_model(file_path):
                self.model_status_label.setText("● 模型就绪")
                self.model_status_label.setStyleSheet("color: #38a169;")
                self.model_info_label.setText(f"{self.translator.get_model_info()}")
                QMessageBox.information(self, "成功", "模型加载成功")
            else:
                self.model_status_label.setText("● 模型加载失败")
                self.model_status_label.setStyleSheet("color: #e53e3e;")
                QMessageBox.critical(self, "错误", "模型加载失败")
    
    def add_term(self):
        """添加术语定义"""
        term_text = self.term_input.text().strip()
        if term_text:
            if "=" in term_text:
                self.terms.append(term_text)
                self.terms_list.addItem(term_text)
                self.term_input.clear()
            else:
                QMessageBox.warning(self, "警告", "术语格式不正确，请使用 'A = B' 格式")
    
    def toggle_advanced_panel(self, checked):
        """切换高级辅助面板的显示/隐藏状态"""
        self.advanced_content.setVisible(checked)
        self.advanced_toggle.setText("▲" if checked else "▼")
    
    def show_term_context_menu(self, pos):
        """显示术语项右键菜单"""
        from PySide6.QtWidgets import QMenu
        menu = QMenu()
        delete_action = menu.addAction("删除")
        action = menu.exec_(self.terms_list.mapToGlobal(pos))
        if action == delete_action:
            selected_items = self.terms_list.selectedItems()
            for item in selected_items:
                self.terms.remove(item.text())
                self.terms_list.takeItem(self.terms_list.row(item))
    
    def import_word(self):
        """导入Word文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "导入Word文件", "", "Word文件 (*.docx)"
        )
        if file_path:
            try:
                text = self.word_handler.import_docx(file_path)
                self.source_text.setText(text)
                QMessageBox.information(self, "成功", "Word文件导入成功")
            except Exception as e:
                logger.error(f"导入Word文件失败: {e}")
                QMessageBox.critical(self, "错误", f"Word文件导入失败: {str(e)}")
    
    def export_word(self):
        """导出Word文件"""
        result_text = self.result_text.toPlainText()
        if not result_text.strip():
            QMessageBox.warning(self, "警告", "译文结果为空，无法导出")
            return
        
        file_path, _ = QFileDialog.getSaveFileName(
            self, "导出Word文件", "", "Word文件 (*.docx)"
        )
        if file_path:
            try:
                self.word_handler.export_docx(file_path, result_text)
                QMessageBox.information(self, "成功", "Word文件导出成功")
            except Exception as e:
                logger.error(f"导出Word文件失败: {e}")
                QMessageBox.critical(self, "错误", f"Word文件导出失败: {str(e)}")
    
    def copy_result(self):
        """复制译文结果"""
        result_text = self.result_text.toPlainText()
        if result_text.strip():
            from PySide6.QtGui import QClipboard
            clipboard = QApplication.clipboard()
            clipboard.setText(result_text)
            QMessageBox.information(self, "成功", "译文已复制到剪贴板")
        else:
            QMessageBox.warning(self, "警告", "译文结果为空，无法复制")
    
    def start_translation(self):
        """开始翻译"""
        source_text = self.source_text.toPlainText().strip()
        if not source_text:
            QMessageBox.warning(self, "警告", "原文内容为空，无法翻译")
            return
        
        if not hasattr(self.translator, 'llama_cpp_available') or not self.translator.llama_cpp_available:
            QMessageBox.warning(self, "警告", "llama-cpp-python库未安装，无法执行翻译功能")
            return
        
        if not self.translator.is_ready:
            QMessageBox.warning(self, "警告", "模型未就绪，无法翻译")
            return
        
        # 禁用翻译按钮
        self.translate_btn.setEnabled(False)
        self.translate_btn.setText("翻译中...")
        
        # 获取上下文和术语
        context = self.context_edit.toPlainText().strip()
        terms = self.terms if self.terms else None
        
        # 执行翻译（在后台线程中执行，避免UI卡顿）
        from PySide6.QtCore import QThread, Signal
        
        class TranslationThread(QThread):
            finished = Signal(str)
            
            def __init__(self, translator, text, context, terms):
                super().__init__()
                self.translator = translator
                self.text = text
                self.context = context
                self.terms = terms
            
            def run(self):
                result = self.translator.translate(self.text, self.context, self.terms)
                self.finished.emit(result)
        
        self.translation_thread = TranslationThread(self.translator, source_text, context, terms)
        self.translation_thread.finished.connect(self.on_translation_finished)
        self.translation_thread.start()
    
    def on_translation_finished(self, result):
        """翻译完成回调"""
        # 启用翻译按钮
        self.translate_btn.setEnabled(True)
        self.translate_btn.setText("开始翻译")
        
        if result:
            self.result_text.setText(result)
        else:
            QMessageBox.critical(self, "错误", "翻译失败，请检查日志")

# 修复导入问题
from PySide6.QtWidgets import QApplication