import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                             QTreeWidget, QTreeWidgetItem, QTabWidget, QLabel, QPushButton, 
                             QFileDialog, QLineEdit, QFormLayout, QSpinBox, QComboBox, QCheckBox,
                             QTextBrowser, QSplitter, QRubberBand, QDoubleSpinBox, QListWidget,
                             QListWidgetItem, QTableWidget, QTableWidgetItem, QHeaderView, QGroupBox, QMenu, QMessageBox, QProgressDialog, QProgressBar)  # 添加 QListWidgetItem 和 QHeaderView
from PyQt5.QtGui import QIcon, QPixmap, QPainter, QPen, QColor, QGuiApplication, QCursor
from PyQt5.QtCore import Qt, QTimer, QRect, QPoint, pyqtSignal, QSize, QThread
import tools_library as tl
import traceback
import datetime
import yaml

class ScreenshotOverlay(QWidget):
    screenshot_taken = pyqtSignal(QPoint, QPoint)

    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint | Qt.Tool)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setStyleSheet("background-color: rgba(0, 0, 0, 1);")  # 非常轻微的暗色
        self.showFullScreen()
        self.setCursor(Qt.CrossCursor)
        self.start = QPoint()
        self.end = QPoint()
        self.rubberBand = QRubberBand(QRubberBand.Rectangle, self)

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.fillRect(self.rect(), QColor(0, 0, 0, 1))  # 几乎透明的背景

    def mousePressEvent(self, event):
        self.start = event.pos()
        self.rubberBand.setGeometry(QRect(self.start, QSize()))
        self.rubberBand.show()

    def mouseMoveEvent(self, event):
        self.rubberBand.setGeometry(QRect(self.start, event.pos()).normalized())

    def mouseReleaseEvent(self, event):
        self.end = event.pos()
        self.rubberBand.hide()
        self.screenshot_taken.emit(self.start, self.end)
        self.close()

class ClickableImageLabel(QLabel):
    clicked = pyqtSignal()
    image_dropped = pyqtSignal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.setAcceptDrops(True)

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.clicked.emit()
        super().mousePressEvent(event)

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dropEvent(self, event):
        files = [u.toLocalFile() for u in event.mimeData().urls()]
        for file_path in files:
            if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                self.image_dropped.emit(file_path)
                break

class WorkerThread(QThread):
    progress = pyqtSignal(list, int, int)
    finished = pyqtSignal(str)
    error = pyqtSignal(str)

    def __init__(self, function, *args, **kwargs):
        super().__init__()
        self.function = function
        self.args = args
        self.kwargs = kwargs
        self.is_stopped = False

    def run(self):
        try:
            # 将 progress_callback 和 stop_check 作为参数传递给函数
            self.kwargs['progress_callback'] = self.progress.emit
            self.kwargs['stop_check'] = self.check_if_stopped
            result = self.function(*self.args, **self.kwargs)
            if not self.is_stopped:
                self.finished.emit(result)
        except Exception as e:
            self.error.emit(str(e))

    def stop(self):
        self.is_stopped = True

    def check_if_stopped(self):
        return self.is_stopped

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.config = self.load_config()
        self.initUI()
        self.screenshot_path = None
        self.overlay = None
        self.current_image_path = None
        
        # 检查输入输出文件夹
        self.check_input_output_directories()
        
        # 使用 QTimer 延迟加载模型
        QTimer.singleShot(100, self.load_model)
        
    def handle_dropped_image(self, file_path):
        self.file_input.setText(file_path)
        self.display_image(file_path)
        self.recognize_single_image()

    def check_input_output_directories(self):
        warnings = []
        
        # 检查标题提取分类
        title_config = self.config.get('title_extraction', {})
        if title_config.get('input_dir') == title_config.get('output_dir') and title_config.get('input_dir'):
            warnings.append("标题提取分类的输入和输出文件夹相同")
        
        # 检查内容提取分类
        content_config = self.config.get('content_keyword_extraction', {})
        if content_config.get('input_dir') == content_config.get('output_dir') and content_config.get('input_dir'):
            warnings.append("内容提取分类的输入和输出文件夹相同")
        
        # 检查图片信息提取
        key_info_config = self.config.get('key_info_extraction', {})
        if key_info_config.get('input_dir') == key_info_config.get('output_dir') and key_info_config.get('input_dir'):
            warnings.append("图片信息提取的输入和输出文件夹相同")
        
        # 如果有警告，显示弹窗
        if warnings:
            warning_message = "\n".join(warnings)
            QMessageBox.warning(self, "输入输出文件夹警告", 
                                f"以下工具的输入和输出文件夹相同，可能会导致原始文件被覆盖：\n\n{warning_message}\n\n请在使用前修改相关设置。")

    def load_config(self):
        try:
            with open('config.yaml', 'r', encoding='utf-8') as file:
                return yaml.safe_load(file)
        except FileNotFoundError:
            print("配置文件不存在，将使用默认配置")
            return {}
        except yaml.YAMLError as e:
            print(f"读取配置文件时出错: {e}")
            return {}

    def save_config(self):
        with open('config.yaml', 'w', encoding='utf-8') as file:
            yaml.dump(self.config, file, allow_unicode=True)

    def initUI(self):
        self.setWindowTitle('ocr-classify')
        self.setGeometry(100, 100, 1400, 900)
        self.setWindowIcon(QIcon(r'images\logo.png'))

        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        main_layout = QHBoxLayout()
        central_widget.setLayout(main_layout)

        # 创建一个 QSplitter
        splitter = QSplitter(Qt.Horizontal)
        main_layout.addWidget(splitter)

        # 左侧树形控件
        self.tree = QTreeWidget()
        self.tree.setHeaderLabel('工具列表')
        splitter.addWidget(self.tree)

        # 添加点
        root = QTreeWidgetItem(self.tree, ['OCR工具列表'])
        tools = [
            '单张图片识别',
            '标题提取分类',
            '内容提取分类',
            '图片信息提取',
            '模型参数设置',
            '帮助'  # 新添加的选项
        ]
        for tool in tools:
            QTreeWidgetItem(root, [tool])

        self.tree.expandAll()
        self.tree.itemClicked.connect(self.on_item_clicked)

        # 右侧标签页
        self.tab_widget = QTabWidget()
        self.tab_widget.setTabsClosable(True)
        self.tab_widget.tabCloseRequested.connect(self.close_tab)
        splitter.addWidget(self.tab_widget)

        # 设置初始的分割比例
        splitter.setSizes([300, 1100])  # 左侧宽度300，右侧为1100

    def load_ui_state(self):
        # 这里可以放置一些通用的加载逻辑
        pass

    def save_ui_state(self):
        # 保存单张图片识别参数
        if hasattr(self, 'file_input') and hasattr(self, 'visualize_checkbox'):
            self.config['single_image_recognition'] = {
                'last_input_path': self.file_input.text(),
                'visualize': self.visualize_checkbox.isChecked()
            }

        # 保存标题提取分类参数
        if all(hasattr(self, attr) for attr in ['input_dir', 'output_dir', 'keyword_list', 'start_h', 'end_h', 'base_width']):
            self.config['title_extraction'] = {
                'input_dir': self.input_dir.text(),
                'output_dir': self.output_dir.text(),
                'keywords': [self.keyword_list.item(i).text() for i in range(self.keyword_list.count())],
                'start_h': self.start_h.value(),
                'end_h': self.end_h.value(),
                'base_width': self.base_width.value()
            }

        # 保存内容提取分类参数
        if all(hasattr(self, attr) for attr in ['content_input_dir', 'content_output_dir', 'content_keyword_list', 'content_start_h', 'content_end_h', 'content_base_width']):
            self.config['content_keyword_extraction'] = {
                'input_dir': self.content_input_dir.text(),
                'output_dir': self.content_output_dir.text(),
                'keywords': [self.content_keyword_list.item(i).text() for i in range(self.content_keyword_list.count())],
                'start_h': self.content_start_h.value(),
                'end_h': self.content_end_h.value(),
                'base_width': self.content_base_width.value()
            }

        # 保存图片信息提取参数
        if all(hasattr(self, attr) for attr in ['key_info_input_dir', 'key_info_output_dir', 'key_info_start_h', 'key_info_end_h', 'key_info_base_width', 'key_info_table']):
            self.config['key_info_extraction'] = {
                'input_dir': self.key_info_input_dir.text(),
                'output_dir': self.key_info_output_dir.text(),
                'start_h': self.key_info_start_h.value(),
                'end_h': self.key_info_end_h.value(),
                'base_width': self.key_info_base_width.value(),
                'key_info_params': self.get_key_info_params()
            }

        

        self.save_config()

    def closeEvent(self, event):
        self.save_ui_state()
        super().closeEvent(event)

    def _execute_extraction(self, input_widget, output_widget, keyword_list_widget,
                        start_h_widget, end_h_widget, base_width_widget,
                        result_table, extraction_function, operation_name,
                        post_process_function=None):
        input_dir = input_widget.text()
        output_dir = output_widget.text()
        keywords = [keyword_list_widget.item(i).text() for i in range(keyword_list_widget.count())]
        start_h = start_h_widget.value()
        end_h = end_h_widget.value()
        base_width = base_width_widget.value()

        if not input_dir or not output_dir:
            result_table.setRowCount(1)
            result_table.setItem(0, 0, QTableWidgetItem("请选择输入和出文件夹"))
            return

        result_table.setRowCount(0)

        def progress_callback(csv_text, current, total):
            row = result_table.rowCount()
            result_table.insertRow(row)
            for col, text in enumerate(csv_text):
                if col == 1:  # 第二为定位图片路径
                    item = QTableWidgetItem("点击打开")
                    item.setData(Qt.UserRole, text)  # 存储图片路径
                    result_table.setItem(row, col, item)
                else:
                    result_table.setItem(row, col, QTableWidgetItem(str(text)))
            result_table.resizeRowToContents(row)
            
            # 调整列的宽度
            for i in range(3):
                result_table.resizeColumnToContents(i)
            
            QApplication.processEvents()

        try:
            # 创建一个字典来存储所有参数
            params = {
                'keywords': keywords,
                'progress_callback': progress_callback,
                'start_h': start_h,
                'end_h': end_h,
                'base_width': base_width
            }
            
            # 获取 extraction_function 的参数列表
            import inspect
            func_params = inspect.signature(extraction_function).parameters
            
            # 只传递 extraction_function 接受的参数
            filtered_params = {k: v for k, v in params.items() if k in func_params}
            
            # 始终传递 input_dir 和 output_dir 作为位置数
            output_folder = extraction_function(input_dir, output_dir, **filtered_params)
            
            if output_folder is None:
                raise ValueError(f"{operation_name}函数返回了 None")
            if post_process_function:
                post_process_function(output_folder, keywords)
            self.show_result_dialog(f"{operation_name}完成", output_folder)
        except Exception as e:
            self.show_result_dialog(f"执行过程中出错: {str(e)}", output_dir)
            print(f"详细错误信息: {traceback.format_exc()}")

    def load_model(self):
        try:
            # 从配置文件中读取模型参数
            if 'model_params' in self.config:
                params = self.config['model_params']
            else:
                # 如果配置文件中没有模型参数，使用默认参数
                params = {
                    'det_model': 'models/ch_PP-OCRv3_det_infer',
                    'cls_model': 'models/ch_ppocr_mobile_v2.0_cls_infer',
                    'rec_model': 'models/ch_PP-OCRv3_rec_infer',
                    'rec_label_file': 'models/ppocr_keys_v1.txt',
                    'device': 'cpu',
                    'backend': 'default',
                    'device_id': 0,
                    'cpu_thread_num': 9,
                    'cls_bs': 1,
                    'rec_bs': 6,
                    'use_cls_model': False
                    
                }# 'process_num': 5,# 'thread_num': 1,

            # 使用参初始化模型
            success = tl.initialize_model(params)
            if success:
                print("模型加载完成，使用的参数：", params)
            else:
                print("模型加载失败")
        except Exception as e:
            print(f"模型加载失败: {str(e)}")

    def on_item_clicked(self, item, column):
        tool_name = item.text(0)
        if tool_name == '单张图片识别':
            self.create_or_switch_to_single_image_tab()
        elif tool_name == '标题提取分类':
            self.create_or_switch_to_title_extraction_tab()
        elif tool_name == '内容提取分类':
            self.create_or_switch_to_content_keyword_tab()
        elif tool_name == '图片信息提取':
            self.create_or_switch_to_key_info_extraction_tab()
        elif tool_name == '模型参数设置':
            self.create_or_switch_to_model_params_tab()
        elif tool_name == '帮助':  # 新添加的条件
            self.create_or_switch_to_help_tab()

    def create_or_switch_to_single_image_tab(self):
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '单张图片识别':
                self.tab_widget.setCurrentIndex(i)
                return
        self.create_single_image_tab()

    def create_or_switch_to_title_extraction_tab(self):
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '标题提取分类':
                self.tab_widget.setCurrentIndex(i)
                return
        self.create_title_extraction_tab()

    def create_or_switch_to_content_keyword_tab(self):
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '内容提取分类':
                self.tab_widget.setCurrentIndex(i)
                return
        self.create_content_keyword_tab()

    def create_or_switch_to_key_info_extraction_tab(self):
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '图片信息提取':
                self.tab_widget.setCurrentIndex(i)
                return
        self.create_key_info_extraction_tab()

    def create_or_switch_to_model_params_tab(self):
        # 检查是否已存在模型参数设置标签页
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '模型参数设置':
                self.tab_widget.setCurrentIndex(i)
                return
        
        # 如果不存在，则创建新的标签
        self.create_model_params_tab()

    def create_or_switch_to_help_tab(self):
        # 检查是否已存在帮助标签页
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '帮助':
                self.tab_widget.setCurrentIndex(i)
                return

        # 如果不存在，创建新的帮助标签页
        help_tab = QWidget()
        layout = QVBoxLayout()
        help_tab.setLayout(layout)

        # 创建文本浏览器来显示帮助内容
        text_browser = QTextBrowser()
        text_browser.setOpenExternalLinks(True)  # 允许打开外部链接
        layout.addWidget(text_browser)

        # 获取当前脚本的目录
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 构建图片的绝对路径
        image_path = os.path.join(current_dir, 'images', 'qq.jpg')
        
        # 直接在代码中定义帮助内容，使用HTML格式
        help_content = f"""
        <div style="margin-bottom: 30px;">
            <h2>项目地址：<a href="https://gitee.com/sevenlusir/ocr-classify">https://gitee.com/sevenlusir/ocr-classify</a></h2>
        </div>
        <div>
            <img src="{image_path}" alt="QQ群二维码" style="width:300px;">
        </div>
        """

        # 设置帮助内容
        text_browser.setHtml(help_content)

        # 添加新的标签页
        self.tab_widget.addTab(help_tab, '帮助')
        self.tab_widget.setCurrentWidget(help_tab)

    def close_tab(self, index):
        self.tab_widget.removeTab(index)

    def create_single_image_tab(self):
        # 检查是否已存在相同的签页
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '单张图片识别':
                self.tab_widget.setCurrentIndex(i)
                return

        # 如果不存在，创建新的标签页
        tab = QWidget()
        tab.setObjectName('single_image_tab')
        layout = QVBoxLayout()
        tab.setLayout(layout)

        # 第一排控件
        top_layout = QHBoxLayout()
        self.file_input = QLineEdit()
        browse_button = QPushButton('浏览')
        browse_button.clicked.connect(self.browse_file)
        screenshot_button = QPushButton('屏幕截图')
        screenshot_button.clicked.connect(self.take_screenshot)
        self.visualize_checkbox = QCheckBox('显示可视化结果')
        self.visualize_checkbox.setChecked(False)  # 默认选中
        top_layout.addWidget(self.file_input)
        top_layout.addWidget(browse_button)
        top_layout.addWidget(screenshot_button)
        top_layout.addWidget(self.visualize_checkbox)
        layout.addLayout(top_layout)

        # 第二排控件
        splitter = QSplitter(Qt.Horizontal)
        self.image_label = ClickableImageLabel("拖拽图片到此区域进行识别")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setStyleSheet("border: 1px solid black;")
        self.image_label.setCursor(QCursor(Qt.PointingHandCursor))
        self.image_label.clicked.connect(self.open_image)
        self.image_label.image_dropped.connect(self.handle_dropped_image)
        
        # 创建一个新的 widget 来包含结果文本
        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        
        # 创建一个新的 QTextBrowser 来包含结果文本
        self.result_text = QTextBrowser()
        self.result_text.setOpenExternalLinks(False)  # 禁止打开外部链接
        right_layout.addWidget(self.result_text)
        
        # 创建复制按钮并添加到结果文本下方
        self.copy_button = QPushButton('复制结果')
        self.copy_button.clicked.connect(self.copy_result)
        right_layout.addWidget(self.copy_button)
        
        splitter.addWidget(self.image_label)
        splitter.addWidget(right_widget)
        splitter.setSizes([840, 560])  # 60% and 40% of 1400
        layout.addWidget(splitter)

        self.tab_widget.addTab(tab, '单张图片识别')
        self.tab_widget.setCurrentWidget(tab)

        # 从配置中读取参数
        single_image_config = self.config.get('single_image_recognition', {})
        self.file_input.setText(single_image_config.get('last_input_path', ''))
        self.visualize_checkbox.setChecked(single_image_config.get('visualize', True))

    def browse_file(self):
        file_name, _ = QFileDialog.getOpenFileName(self, '选择图', '', 'Image Files (*.png *.jpg *.bmp)')
        if file_name:
            self.file_input.setText(file_name)
            self.display_image(file_name)
            self.recognize_single_image()  # 自动进行 OCR 识别

    def take_screenshot(self):
        if self.overlay is None or not self.overlay.isVisible():
            self.overlay = ScreenshotOverlay()
            self.overlay.screenshot_taken.connect(self.capture_screen)
            self.hide()
            QTimer.singleShot(100, self.show_overlay)
        else:
            print("Screenshot overlay is already visible")

    def show_overlay(self):
        if self.overlay:
            self.overlay.showFullScreen()
        else:
            print("Failed to create overlay")
            self.show()  # 如果无法创建覆盖层，则显示主窗口

    def capture_screen(self, start, end):
        if self.overlay:
            self.overlay.hide()
        QTimer.singleShot(100, lambda: self._do_capture_screen(start, end))

    def _do_capture_screen(self, start, end):
        try:
            screen = QApplication.primaryScreen()
            original_pixmap = screen.grabWindow(0)
            
            # 计算选定区
            x = min(start.x(), end.x())
            y = min(start.y(), end.y())
            width = abs(start.x() - end.x())
            height = abs(start.y() - end.y())
            
            # 裁剪图片
            cropped_pixmap = original_pixmap.copy(x, y, width, height)
            
            if not os.path.exists('screenshots'):
                os.makedirs('screenshots')
            
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            self.screenshot_path = os.path.join('screenshots', f'screenshot_{timestamp}.png')
            
            cropped_pixmap.save(self.screenshot_path, 'png')
            print(f"Screenshot saved to: {self.screenshot_path}")
            
            self.show()  # 显示主窗口
            self.display_image(self.screenshot_path)
            self.file_input.setText(self.screenshot_path)
            
            # 直 OCR 识别，不使用 QTimer
            self.recognize_single_image()
        except Exception as e:
            print(f"截图出错: {e}")
            print(traceback.format_exc())

    def display_image(self, image_path):
        if image_path and os.path.exists(image_path):
            pixmap = QPixmap(image_path)
            if not pixmap.isNull():
                scaled_pixmap = pixmap.scaled(self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
                self.image_label.setPixmap(scaled_pixmap)
                self.image_label.setAlignment(Qt.AlignCenter)
                self.current_image_path = image_path
            else:
                print(f"无法加载图像: {image_path}")
        else:
            print(f"图像文件不存在: {image_path}")

    def open_image(self):
        if self.current_image_path and os.path.exists(self.current_image_path):
            os.startfile(self.current_image_path)
        else:
            print("没有可打开的图像")

    def recognize_single_image(self):
        if not tl.ppocr_v3:
            self.result_text.setHtml("模型尚未加载完成，正在尝试加载...")
            self.load_model()
            if not tl.ppocr_v3:
                self.result_text.setHtml("模型加载失败，请检查模参数设置")
                return

        image_path = self.file_input.text()
        if not image_path:
            print("No image path provided")
            return
        self.result_text.setHtml("正在进行OCR识别...")
        QApplication.processEvents()  # 刷新UI
        try:
            visualize = self.visualize_checkbox.isChecked()
            result, vis_path = tl.single_image_ocr(image_path, visualize)
            print(f"OCR Result: {result}")  # 调试输出
            if result:
                # 将列表转换为字符串
                result_text = "<br>".join(result)
                self.result_text.setHtml(result_text)
                # 根据复选框状态决定显示原图还是可视化结果
                if visualize and vis_path:
                    self.display_image(vis_path)
                else:
                    self.display_image(image_path)
            else:
                self.result_text.setHtml("OCR识别未返回结果")
        except Exception as e:
            error_message = f"OCR识别出错: {str(e)}"
            print(error_message)
            self.result_text.setHtml(error_message)
        finally:
            QApplication.processEvents()  # 再次刷新UI，确保结果显示

    def copy_result(self):
        clipboard = QApplication.clipboard()
        clipboard.setText(self.result_text.toPlainText())
        self.copy_button.setText('已复制')
        QTimer.singleShot(2000, lambda: self.copy_button.setText('复制'))

    def create_title_extraction_tab(self):
        tab = QWidget()
        main_layout = QVBoxLayout()
        tab.setLayout(main_layout)

        # 创建上部参数区域
        params_widget = QWidget()
        params_layout = QFormLayout()
        params_widget.setLayout(params_layout)

        # 输入文件夹路径
        self.input_dir = QLineEdit()
        browse_input_button = QPushButton("浏览")
        browse_input_button.clicked.connect(self.browse_title_input_directory)
        input_layout = QHBoxLayout()
        input_layout.addWidget(self.input_dir)
        input_layout.addWidget(browse_input_button)
        params_layout.addRow("输入文件夹:", input_layout)

        # 输出文件夹路径
        self.output_dir = QLineEdit()
        browse_output_button = QPushButton("浏览")
        browse_output_button.clicked.connect(self.browse_title_output_directory)
        output_layout = QHBoxLayout()
        output_layout.addWidget(self.output_dir)
        output_layout.addWidget(browse_output_button)
        params_layout.addRow("输出文件夹:", output_layout)

        # 裁剪参数
        self.start_h = QDoubleSpinBox()
        self.start_h.setRange(0, 1)
        self.start_h.setSingleStep(0.01)
        self.start_h.setValue(0.05)
        params_layout.addRow("裁剪开始高度比例:", self.start_h)

        self.end_h = QDoubleSpinBox()
        self.end_h.setRange(0, 1)
        self.end_h.setSingleStep(0.01)
        self.end_h.setValue(0.35)
        params_layout.addRow("裁剪结束高度比例:", self.end_h)

        self.base_width = QSpinBox()
        self.base_width.setRange(0, 10000)
        self.base_width.setValue(960)
        params_layout.addRow("基础宽度:", self.base_width)

        # 将参数区域添加到主布局
        main_layout.addWidget(params_widget)

        # 创建关键词列表区域
        keyword_widget = QWidget()
        keyword_layout = QVBoxLayout()
        keyword_widget.setLayout(keyword_layout)

        self.keyword_list = QListWidget()
        keyword_layout.addWidget(QLabel("标题关键词列表:"))
        keyword_layout.addWidget(self.keyword_list)

        # 添加和删除关键词的按钮
        button_layout = QHBoxLayout()
        self.new_keyword = QLineEdit()
        add_keyword_button = QPushButton("添加关键词")
        add_keyword_button.clicked.connect(self.add_keyword)
        delete_keyword_button = QPushButton("删除选中关键词")
        delete_keyword_button.clicked.connect(self.delete_keywords)
        button_layout.addWidget(self.new_keyword)
        button_layout.addWidget(add_keyword_button)
        button_layout.addWidget(delete_keyword_button)
        keyword_layout.addLayout(button_layout)

        # 将关键词列表区域添加到主布局，并设置为可扩展
        main_layout.addWidget(keyword_widget, 1)  # 1表示可以扩展

        # 添加执行和停止按钮
        execute_layout = QHBoxLayout()
        self.title_execute_button = QPushButton("运行")
        self.title_execute_button.clicked.connect(self.execute_title_extraction)
        self.title_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.title_stop_button = QPushButton("取消")
        self.title_stop_button.clicked.connect(self.stop_title_extraction)
        self.title_stop_button.setEnabled(False)  # 初始状态设置不可用
        execute_layout.addWidget(self.title_execute_button)
        execute_layout.addWidget(self.title_stop_button)
        main_layout.addLayout(execute_layout)

        # 创建进度条和标签的布局
        progress_layout = QHBoxLayout()

        # 添加进度条和标签
        self.title_progress_bar = QProgressBar()
        self.title_progress_label = QLabel("准备开始...")
        progress_layout.addWidget(self.title_progress_bar)
        progress_layout.addWidget(self.title_progress_label)

        # 将进度条布局添加到主布局
        main_layout.addLayout(progress_layout)

        self.tab_widget.addTab(tab, '标题提取分类')
        self.tab_widget.setCurrentWidget(tab)

        # 从配置中读取参数
        title_extraction_config = self.config.get('title_extraction', {})
        self.input_dir.setText(title_extraction_config.get('input_dir', ''))
        self.output_dir.setText(title_extraction_config.get('output_dir', ''))
        self.keyword_list.clear()
        self.keyword_list.addItems(title_extraction_config.get('keywords', []))
        self.start_h.setValue(title_extraction_config.get('start_h', 0.05))
        self.end_h.setValue(title_extraction_config.get('end_h', 0.35))
        self.base_width.setValue(title_extraction_config.get('base_width', 960))

    def browse_title_input_directory(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择输入文件夹")
        if dir_path:
            if dir_path == self.output_dir.text():
                QMessageBox.warning(self, "警告", "输入文件夹不能与输出文件夹相同，请重新选择")
            else:
                self.input_dir.setText(dir_path)
                self.save_ui_state()

    def browse_title_output_directory(self):
        self._browse_output_directory(self.input_dir, self.output_dir)
        self.save_ui_state()

    def add_keyword(self):
        keyword = self.new_keyword.text().strip()
        if keyword:
            self.keyword_list.addItem(keyword)
            self.new_keyword.clear()
            self.save_ui_state()

    def delete_keywords(self):
        for item in self.keyword_list.selectedItems():
            self.keyword_list.takeItem(self.keyword_list.row(item))
        self.save_ui_state()

    def execute_title_extraction(self):
        input_dir = self.input_dir.text()
        output_dir = self.output_dir.text()
        start_h = self.start_h.value()
        end_h = self.end_h.value()
        base_width = self.base_width.value()
        keywords = [self.keyword_list.item(i).text() for i in range(self.keyword_list.count())]

        if not keywords:
            QMessageBox.warning(self, "警告", "关键词列表为空，请添加至少一个关键词")
            return

        if not input_dir or not output_dir:
            self.title_progress_label.setText("请选择输入和输出文件夹")
            return

        if input_dir == output_dir:
            QMessageBox.warning(self, "警告", "输出文件夹不能与输入文件夹相同，请重新选择输出文件夹。")
            return

        # 重置进度条和标签
        self.title_progress_bar.setValue(0)
        self.title_progress_label.setText("正在处理图片...")

        # 更新按钮状态
        self.title_execute_button.setEnabled(False)
        self.title_execute_button.setStyleSheet("background-color: #cccccc; color: #666666;")
        self.title_stop_button.setEnabled(True)
        self.title_stop_button.setStyleSheet("background-color: #f44336; color: white;")

        # 创建并启动工作线程
        self.title_worker = WorkerThread(tl.title_keyword_categorization, 
                                           input_dir, output_dir, keywords, start_h=start_h, 
                                           end_h=end_h, base_width=base_width)
        self.title_worker.progress.connect(self.update_title_progress)
        self.title_worker.finished.connect(self.on_title_extraction_finished)
        self.title_worker.error.connect(self.on_title_extraction_error)
        self.title_worker.start()

    def update_title_progress(self, csv_text, current, total):
        if not self.title_worker.is_stopped:
            progress = int((current / total) * 100)
            self.title_progress_bar.setValue(progress)
            self.title_progress_label.setText(f"正在处理: {current}/{total}")

    def on_title_extraction_finished(self, output_folder):
        self.title_progress_bar.setValue(100)
        self.title_progress_label.setText("处理完成")
        self.show_result_dialog("标题提取分类完成", output_folder)
        
        # 重置按钮状态
        self.title_execute_button.setEnabled(True)
        self.title_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.title_stop_button.setEnabled(False)
        self.title_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")

    def on_title_extraction_error(self, error_message):
        self.title_progress_label.setText("处理出错")
        self.show_result_dialog(f"执行过程中出错: {error_message}", self.output_dir.text())
        self.title_execute_button.setEnabled(True)
        self.title_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色
        self.title_stop_button.setEnabled(False)
        self.title_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色

    def stop_title_extraction(self):
        if hasattr(self, 'title_worker') and self.title_worker.isRunning():
            self.title_worker.stop()
            self.title_progress_bar.setValue(0)  # 重置进度条
            self.title_progress_label.setText("已取消")  # 更新标签
            self.title_stop_button.setEnabled(False)
            self.title_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色
            self.title_execute_button.setEnabled(True)
            self.title_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色

    
    def create_content_keyword_tab(self):
        tab = QWidget()
        main_layout = QVBoxLayout()
        tab.setLayout(main_layout)

        # 创建上部参数区域
        params_widget = QWidget()
        params_layout = QFormLayout()
        params_widget.setLayout(params_layout)

        # 输入文件夹路径
        self.content_input_dir = QLineEdit()
        browse_input_button = QPushButton("浏览")
        browse_input_button.clicked.connect(self.browse_content_input_directory)
        input_layout = QHBoxLayout()
        input_layout.addWidget(self.content_input_dir)
        input_layout.addWidget(browse_input_button)
        params_layout.addRow("输入文件夹:", input_layout)

        # 输出文件夹路径
        self.content_output_dir = QLineEdit()
        browse_output_button = QPushButton("浏览")
        browse_output_button.clicked.connect(self.browse_content_output_directory)
        output_layout = QHBoxLayout()
        output_layout.addWidget(self.content_output_dir)
        output_layout.addWidget(browse_output_button)
        params_layout.addRow("输出文件夹:", output_layout)

        # 裁剪参数
        self.content_start_h = QDoubleSpinBox()
        self.content_start_h.setRange(0, 1)
        self.content_start_h.setSingleStep(0.01)
        params_layout.addRow("裁剪开始高度比例:", self.content_start_h)

        self.content_end_h = QDoubleSpinBox()
        self.content_end_h.setRange(0, 1)
        self.content_end_h.setSingleStep(0.01)
        params_layout.addRow("裁剪结束高度比例:", self.content_end_h)

        self.content_base_width = QSpinBox()
        self.content_base_width.setRange(0, 10000)
        params_layout.addRow("基础宽度:", self.content_base_width)

        # 将参数区域添加到主布局
        main_layout.addWidget(params_widget)

        # 创建关键词列表区域
        keyword_widget = QWidget()
        keyword_layout = QVBoxLayout()
        keyword_widget.setLayout(keyword_layout)

        self.content_keyword_list = QListWidget()
        keyword_layout.addWidget(QLabel("内容关键词列表:"))
        keyword_layout.addWidget(self.content_keyword_list)

        # 添加和删除关键词的按钮
        button_layout = QHBoxLayout()
        self.new_content_keyword = QLineEdit()
        add_keyword_button = QPushButton("添加关键词")
        add_keyword_button.clicked.connect(self.add_content_keyword)
        delete_keyword_button = QPushButton("删除选中关键词")
        delete_keyword_button.clicked.connect(self.delete_content_keywords)
        button_layout.addWidget(self.new_content_keyword)
        button_layout.addWidget(add_keyword_button)
        button_layout.addWidget(delete_keyword_button)
        keyword_layout.addLayout(button_layout)

        # 将关键词列表区域添加到主布局，并设置为可扩展
        main_layout.addWidget(keyword_widget, 1)  # 1表示可以扩展

        # 添加执行和停止按钮
        execute_layout = QHBoxLayout()
        self.content_execute_button = QPushButton("运行")
        self.content_execute_button.clicked.connect(self.execute_content_keyword_extraction)
        self.content_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.content_stop_button = QPushButton("取消")
        self.content_stop_button.clicked.connect(self.stop_content_keyword_extraction)
        # self.content_stop_button.setStyleSheet("background-color: #f44336; color: white;")
        self.content_stop_button.setEnabled(False)  # 初始状态设置为不可用
        execute_layout.addWidget(self.content_execute_button)
        execute_layout.addWidget(self.content_stop_button)
        main_layout.addLayout(execute_layout)

        # 创建进度条和标签的布局
        progress_layout = QHBoxLayout()

        # 添加进度条和标签
        self.content_progress_bar = QProgressBar()
        self.content_progress_label = QLabel("准备开始...")
        progress_layout.addWidget(self.content_progress_bar)
        progress_layout.addWidget(self.content_progress_label)

        # 将进度条布局添加到主布局
        main_layout.addLayout(progress_layout)

        self.tab_widget.addTab(tab, '内容提取分类')
        self.tab_widget.setCurrentWidget(tab)

        # 从配置中读取参数
        content_config = self.config.get('content_keyword_extraction', {})
        self.content_input_dir.setText(content_config.get('input_dir', ''))
        self.content_output_dir.setText(content_config.get('output_dir', ''))
        self.content_keyword_list.clear()
        self.content_keyword_list.addItems(content_config.get('keywords', []))
        self.content_start_h.setValue(content_config.get('start_h', 0.0))
        self.content_end_h.setValue(content_config.get('end_h', 0.0))
        self.content_base_width.setValue(content_config.get('base_width', 960))

    def browse_content_input_directory(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择输入文件夹")
        if dir_path:
            if dir_path == self.content_output_dir.text():
                QMessageBox.warning(self, "警告", "输入文件夹不能与输出文件夹相同，请重新选择")
            else:
                self.content_input_dir.setText(dir_path)
                self.save_ui_state()

    def browse_content_output_directory(self):
        self._browse_output_directory(self.content_input_dir, self.content_output_dir)
        self.save_ui_state()

    def add_content_keyword(self):
        keyword = self.new_content_keyword.text().strip()
        if keyword:
            self.content_keyword_list.addItem(keyword)
            self.new_content_keyword.clear()
            self.save_ui_state()

    def delete_content_keywords(self):
        for item in self.content_keyword_list.selectedItems():
            self.content_keyword_list.takeItem(self.content_keyword_list.row(item))
        self.save_ui_state()

    def execute_content_keyword_extraction(self):
        input_dir = self.content_input_dir.text()
        output_dir = self.content_output_dir.text()
        start_h = self.content_start_h.value()
        end_h = self.content_end_h.value()
        base_width = self.content_base_width.value()
        keywords = [self.content_keyword_list.item(i).text() for i in range(self.content_keyword_list.count())]

        if not keywords:
            QMessageBox.warning(self, "警告", "关键词列表为空，请添加至少一个关键词")
            return

        if not input_dir or not output_dir:
            self.content_progress_label.setText("请选择输入和输出文件夹")
            return

        if input_dir == output_dir:
            QMessageBox.warning(self, "警告", "输出文件夹不能与输入文件夹相同，请重新选择输出文件夹。")
            return

        # 重置进度条和标签
        self.content_progress_bar.setValue(0)
        self.content_progress_label.setText("正在处理图片...")

        # 更新按钮状态
        self.content_execute_button.setEnabled(False)
        self.content_execute_button.setStyleSheet("background-color: #cccccc; color: #666666;")
        self.content_stop_button.setEnabled(True)
        self.content_stop_button.setStyleSheet("background-color: #f44336; color: white;")

        # 创建并启动工作线程
        self.content_worker = WorkerThread(tl.content_keyword_categorization_and_classification, 
                                           input_dir, output_dir, keywords, start_h=start_h, 
                                           end_h=end_h, base_width=base_width)
        self.content_worker.progress.connect(self.update_content_progress)
        self.content_worker.finished.connect(self.on_content_extraction_finished)
        self.content_worker.error.connect(self.on_content_extraction_error)
        self.content_worker.start()

    def update_content_progress(self, csv_text, current, total):
        if not self.content_worker.is_stopped:
            progress = int((current / total) * 100)
            self.content_progress_bar.setValue(progress)
            self.content_progress_label.setText(f"正在处理: {current}/{total}")

    def on_content_extraction_finished(self, output_folder):
        self.content_progress_bar.setValue(100)
        self.content_progress_label.setText("处理完成")
        self.show_result_dialog("内容提取分类完成", output_folder)
        
        # 重置按钮状态
        self.content_execute_button.setEnabled(True)
        self.content_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.content_stop_button.setEnabled(False)
        self.content_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")

    def on_content_extraction_error(self, error_message):
        self.content_progress_label.setText("处理出错")
        self.show_result_dialog(f"执行过程中出错: {error_message}", self.content_output_dir.text())
        self.execute_button.setEnabled(True)
        self.execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色
        self.stop_button.setEnabled(False)
        self.stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色

    def on_extraction_error(self, error_message):
        self.progress_label.setText("处理出错")
        self.show_result_dialog(f"执行过程中出错: {error_message}", self.key_info_output_dir.text())
        print(f"详细错误信息: {traceback.format_exc()}")
        
        # 重置按钮状态
        self.execute_button.setEnabled(True)
        self.execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.stop_button.setEnabled(False)
        self.stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")

    def get_key_info_params(self):
        tar_text_dic = {}
        for row in range(self.key_info_table.rowCount()):
            key_item = self.key_info_table.item(row, 0)
            if key_item is None or key_item.text().strip() == "":
                continue  # 跳过空行
            key = key_item.text()
            
            try:
                x_num = float(self.key_info_table.item(row, 1).text()) if self.key_info_table.item(row, 1) else 0.0
                up_num = float(self.key_info_table.item(row, 2).text()) if self.key_info_table.item(row, 2) else 0.0
                do_num = float(self.key_info_table.item(row, 3).text()) if self.key_info_table.item(row, 3) else 0.0
            except ValueError:
                # 如果无法转换为浮点数，使用默认值
                x_num, up_num, do_num = 0.0, 0.0, 0.0
            
            tar_text_dic[key] = {
            'x_num': x_num, 
            'up_num': up_num, 
            'do_num': do_num,
            'find_flag': 1 if self.find_flag_combo.currentText() == "精确查找" else 2
             }
        return tar_text_dic

    def on_key_info_item_changed(self, item):
        # 检查是否所有必要的 UI 元素都已创建
        if all(hasattr(self, attr) for attr in ['key_info_input_dir', 'key_info_output_dir', 'key_info_start_h', 'key_info_end_h', 'key_info_base_width', 'key_info_table']):
            self.save_ui_state()  # 当格项被修改时保存状态

    def on_key_info_crop_changed(self):
        self.save_ui_state()

    def create_model_params_tab(self):
        # 检查是否已存在相同的标签页
        for i in range(self.tab_widget.count()):
            if self.tab_widget.tabText(i) == '模型参数设置':
                self.tab_widget.setCurrentIndex(i)
                return

        # 如果不存在，创建新的标签页
        tab = QWidget()
        layout = QFormLayout()
        tab.setLayout(layout)

        # 从配置文件中读取模型参数
        model_params = self.config.get('model_params', {})

        # 创建输入字段并设置初始值
        self.det_model = QLineEdit(model_params.get('det_model', 'models/ch_PP-OCRv3_det_infer'))
        self.cls_model = QLineEdit(model_params.get('cls_model', 'models/ch_ppocr_mobile_v2.0_cls_infer'))
        self.rec_model = QLineEdit(model_params.get('rec_model', 'models/ch_PP-OCRv3_rec_infer'))
        self.rec_label_file = QLineEdit(model_params.get('rec_label_file', 'models/ppocr_keys_v1.txt'))
        self.device = QComboBox()
        self.device.addItems(['cpu', 'gpu', 'kunlunxin'])
        self.device.setCurrentText(model_params.get('device', 'cpu'))
        self.backend = QLineEdit(model_params.get('backend', 'default'))
        self.device_id = QSpinBox()
        self.device_id.setValue(model_params.get('device_id', 0))
        self.cpu_thread_num = QSpinBox()
        self.cpu_thread_num.setRange(1, 32)
        self.cpu_thread_num.setValue(model_params.get('cpu_thread_num', 9))
        self.cls_bs = QSpinBox()
        self.cls_bs.setValue(model_params.get('cls_bs', 1))
        self.rec_bs = QSpinBox()
        self.rec_bs.setValue(model_params.get('rec_bs', 6))
        self.use_cls_model = QCheckBox()
        self.use_cls_model.setChecked(model_params.get('use_cls_model', False))
        # self.thread_num = QSpinBox()
        # self.thread_num.setValue(model_params.get('thread_num', 1))
        # self.use_multi_process = QCheckBox()
        # self.use_multi_process.setChecked(model_params.get('use_multi_process', True))
        # self.process_num = QSpinBox()
        # self.process_num.setValue(model_params.get('process_num', 5))

        # 添加到布局
        layout.addRow("检测模型路径:", self.det_model)
        layout.addRow("分类模型路径:", self.cls_model)
        layout.addRow("识别模型路径:", self.rec_model)
        layout.addRow("识别字典:", self.rec_label_file)
        layout.addRow("设备:", self.device)
        layout.addRow("后端:", self.backend)
        layout.addRow("显卡ID:", self.device_id)
        layout.addRow("CPU核心数:", self.cpu_thread_num)
        layout.addRow("批次图片数:", self.cls_bs)
        layout.addRow("识别批量大小:", self.rec_bs)
        layout.addRow("竖排文字检测:", self.use_cls_model)
        # layout.addRow("线程数:", self.thread_num)
        # layout.addRow("使用多进程:", self.use_multi_process)
        # layout.addRow("进程数:", self.process_num)

        # 添加保存钮
        save_button = QPushButton("保存设置")
        save_button.clicked.connect(self.save_model_params)
        layout.addRow(save_button)

        self.tab_widget.addTab(tab, '模型参数设置')
        self.tab_widget.setCurrentWidget(tab)

    def handle_anchor_click(self, url):
        if url.toString() == "#copy_button":
            self.copy_result()

    def show_result_dialog(self, message, output_folder):
        msg_box = QMessageBox(self)
        msg_box.setWindowTitle("执行结果")
        msg_box.setText(message)
        msg_box.setStandardButtons(QMessageBox.Ok | QMessageBox.Open)
        msg_box.button(QMessageBox.Open).setText("打开输出文件夹")
        
        # 添加logo
        logo = QPixmap('images/logo.png')
        if not logo.isNull():
            msg_box.setIconPixmap(logo.scaled(64, 64, Qt.KeepAspectRatio, Qt.SmoothTransformation))
        
        result = msg_box.exec_()
        
        if result == QMessageBox.Open:
            os.startfile(output_folder)

    def _browse_output_directory(self, input_widget, output_widget):
        input_dir = input_widget.text()
        if not input_dir:
            QMessageBox.warning(self, "警告", "请先选择输入文件夹")
            return

        dir_path = QFileDialog.getExistingDirectory(self, "选择输出文件夹")
        if dir_path:
            if dir_path == input_dir:
                QMessageBox.warning(self, "警告", "输出文件夹不能与输入文件夹相同，请重新选择")
            else:
                output_widget.setText(dir_path)
                self.save_ui_state()

    def stop_key_info_extraction(self):
        if hasattr(self, 'worker') and self.worker.isRunning():
            self.worker.stop()
            self.progress_bar.setValue(0)  # 重置
            self.progress_label.setText("执行已取消")
            self.stop_button.setEnabled(False)
            self.stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色
            self.execute_button.setEnabled(True)
            self.execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色

    def stop_content_keyword_extraction(self):
        if hasattr(self, 'content_worker') and self.content_worker.isRunning():
            self.content_worker.stop()
            self.content_progress_bar.setValue(0)  # 重置
            self.content_progress_label.setText("已取消")
            self.content_stop_button.setEnabled(False)
            self.content_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色
            self.content_execute_button.setEnabled(True)
            self.content_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色

    def stop_title_extraction(self):
        if hasattr(self, 'title_worker') and self.title_worker.isRunning():
            self.title_worker.stop()
            self.title_progress_bar.setValue(0)  # 重置进度条
            self.title_progress_label.setText("已取消")  # 更新标签
            self.title_stop_button.setEnabled(False)
            self.title_stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")  # 设置为灰色
            self.title_execute_button.setEnabled(True)
            self.title_execute_button.setStyleSheet("background-color: #4CAF50; color: white;")  # 设置为绿色

    def create_key_info_extraction_tab(self):
        tab = QWidget()
        main_layout = QVBoxLayout()
        tab.setLayout(main_layout)

        # 创建上部参数区域
        params_widget = QWidget()
        params_layout = QFormLayout()
        params_widget.setLayout(params_layout)

        # 输入文件夹路径
        self.key_info_input_dir = QLineEdit()
        browse_input_button = QPushButton("浏览")
        browse_input_button.clicked.connect(self.browse_key_info_input_directory)
        input_layout = QHBoxLayout()
        input_layout.addWidget(self.key_info_input_dir)
        input_layout.addWidget(browse_input_button)
        params_layout.addRow("输入文件夹:", input_layout)

        # 输出文件夹路径
        self.key_info_output_dir = QLineEdit()
        browse_output_button = QPushButton("浏览")
        browse_output_button.clicked.connect(self.browse_key_info_output_directory)
        output_layout = QHBoxLayout()
        output_layout.addWidget(self.key_info_output_dir)
        output_layout.addWidget(browse_output_button)
        params_layout.addRow("输出文件夹:", output_layout)

        # 裁剪参数
        self.key_info_start_h = QDoubleSpinBox()
        self.key_info_start_h.setRange(0, 1)
        self.key_info_start_h.setSingleStep(0.01)
        params_layout.addRow("裁剪开始高度比例:", self.key_info_start_h)

        self.key_info_end_h = QDoubleSpinBox()
        self.key_info_end_h.setRange(0, 1)
        self.key_info_end_h.setSingleStep(0.01)
        params_layout.addRow("裁剪结束高度比例:", self.key_info_end_h)

        self.key_info_base_width = QSpinBox()
        self.key_info_base_width.setRange(0, 10000)
        params_layout.addRow("基础宽度:", self.key_info_base_width)

            # 添加查找类型下拉框
        self.find_flag_combo = QComboBox()
        self.find_flag_combo.addItems(["精确查找", "模糊查找"])
        params_layout.addRow("查找类型:", self.find_flag_combo)

        

        # 将参数区域添加到主布局
        main_layout.addWidget(params_widget)

        # 创建关键信息参数表格区域
        table_widget = QWidget()
        table_layout = QVBoxLayout()
        table_widget.setLayout(table_layout)

        # 息参数表格
        self.key_info_table = QTableWidget()
        self.key_info_table.setColumnCount(4)
        self.key_info_table.setHorizontalHeaderLabels(['关键信息', 'x_num', 'up_num', 'do_num'])
        self.key_info_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.key_info_table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.key_info_table.customContextMenuRequested.connect(self.show_context_menu)
        self.key_info_table.itemChanged.connect(self.on_key_info_item_changed)

        table_layout.addWidget(QLabel("关键信息参数:"))
        table_layout.addWidget(self.key_info_table)

        # 添加和删除行的按钮
        button_layout = QHBoxLayout()
        add_row_button = QPushButton("添加行")
        add_row_button.clicked.connect(self.add_key_info_row)
        delete_row_button = QPushButton("删除选中行")
        delete_row_button.clicked.connect(self.delete_key_info_row)
        button_layout.addWidget(add_row_button)
        button_layout.addWidget(delete_row_button)
        table_layout.addLayout(button_layout)

        # 将表格区域添加到主布局，并设置为可扩展
        main_layout.addWidget(table_widget, 1)  # 1表示可以扩展

        # 添加执行和停止按钮
        execute_layout = QHBoxLayout()
        self.execute_button = QPushButton("运行")
        self.execute_button.clicked.connect(self.execute_key_info_extraction)
        self.execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.stop_button = QPushButton("取消")
        self.stop_button.clicked.connect(self.stop_key_info_extraction)
        # self.stop_button.setStyleSheet("background-color: #f44336; color: white;")
        self.stop_button.setEnabled(False)
        execute_layout.addWidget(self.execute_button)
        execute_layout.addWidget(self.stop_button)
        main_layout.addLayout(execute_layout)

        # 创建进度条和标签的布局
        progress_layout = QHBoxLayout()

        # 添加进度条和标签
        self.progress_bar = QProgressBar()
        self.progress_label = QLabel("准备开始...")
        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.progress_label)

        # 将进度条布局添加到主布局
        main_layout.addLayout(progress_layout)

        self.tab_widget.addTab(tab, '图片信息提取')
        self.tab_widget.setCurrentWidget(tab)

        # 从配置中读取参数
        key_info_config = self.config.get('key_info_extraction', {})
        self.key_info_input_dir.setText(key_info_config.get('input_dir', ''))
        self.key_info_output_dir.setText(key_info_config.get('output_dir', ''))
        self.key_info_start_h.setValue(key_info_config.get('start_h', 0.0))
        self.key_info_end_h.setValue(key_info_config.get('end_h', 1.0))
        self.key_info_base_width.setValue(key_info_config.get('base_width', 960))
        self.fill_key_info_table(key_info_config.get('key_info_params', {}))

    def browse_key_info_input_directory(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择输入文件夹")
        if dir_path:
            if dir_path == self.key_info_output_dir.text():
                QMessageBox.warning(self, "警告", "输入文件夹不能与输出文件夹相同，请重新选择")
            else:
                self.key_info_input_dir.setText(dir_path)
                self.save_ui_state()

    def browse_key_info_output_directory(self):
        self._browse_output_directory(self.key_info_input_dir, self.key_info_output_dir)
        self.save_ui_state()

    def add_key_info_row(self):
        row_position = self.key_info_table.rowCount()
        self.key_info_table.insertRow(row_position)
        for col in range(4):
            self.key_info_table.setItem(row_position, col, QTableWidgetItem(""))
        self.save_ui_state()

    def delete_key_info_row(self):
        selected_rows = set(index.row() for index in self.key_info_table.selectedIndexes())
        for row in sorted(selected_rows, reverse=True):
            self.key_info_table.removeRow(row)
        self.save_ui_state()

    def show_context_menu(self, position):
        context_menu = QMenu()
        add_action = context_menu.addAction("添加行")
        delete_action = context_menu.addAction("删除行")

        action = context_menu.exec_(self.key_info_table.mapToGlobal(position))
        if action == add_action:
            self.add_key_info_row()
        elif action == delete_action:
            self.delete_key_info_row()

    def execute_key_info_extraction(self):
        input_dir = self.key_info_input_dir.text()
        output_dir = self.key_info_output_dir.text()
        start_h = self.key_info_start_h.value()
        end_h = self.key_info_end_h.value()
        base_width = self.key_info_base_width.value()
        key_info_params = self.get_key_info_params()

        if not key_info_params:
            QMessageBox.warning(self, "警告", "关键信息参数为空，请添加至少一个关键信息")
            return

        if not input_dir or not output_dir:
            self.progress_label.setText("请选择输入和输出文件夹")
            return

        if input_dir == output_dir:
            QMessageBox.warning(self, "警告", "输出文件夹不能与输入文件夹相同，请重新选择输出文件夹。")
            return

        # 重置进度条和标签
        self.progress_bar.setValue(0)
        self.progress_label.setText("正在处理图片...")

        # 更新按钮状态
        self.execute_button.setEnabled(False)
        self.execute_button.setStyleSheet("background-color: #cccccc; color: #666666;")
        self.stop_button.setEnabled(True)
        self.stop_button.setStyleSheet("background-color: #f44336; color: white;")

        # 创建并启动工作线程
        self.worker = WorkerThread(tl.Key_information_extraction, 
                                   input_dir, output_dir, key_info_params, 
                                   start_h=start_h, end_h=end_h, base_width=base_width)
        self.worker.progress.connect(self.update_key_info_progress)
        self.worker.finished.connect(self.on_key_info_extraction_finished)
        self.worker.error.connect(self.on_extraction_error)
        self.worker.start()

    def update_key_info_progress(self, csv_text, current, total):
        if not self.worker.is_stopped:
            progress = int((current / total) * 100)
            self.progress_bar.setValue(progress)
            self.progress_label.setText(f"正在处理: {current}/{total}")

    def on_key_info_extraction_finished(self, output_folder):
        self.progress_bar.setValue(100)
        self.progress_label.setText("处理完成")
        self.show_result_dialog("图片信息提取完成", output_folder)
        
        # 重置按钮状态
        self.execute_button.setEnabled(True)
        self.execute_button.setStyleSheet("background-color: #4CAF50; color: white;")
        self.stop_button.setEnabled(False)
        self.stop_button.setStyleSheet("background-color: #cccccc; color: #666666;")

    def fill_key_info_table(self, key_info_params):
        if not hasattr(self, 'key_info_table'):
            print("Key info table is not ready yet")
            return

        # 断开信号连接，以防止在填充表格时触发保存
        try:
            self.key_info_table.itemChanged.disconnect()
        except TypeError:
            # 如果信号没有连接，会抛出 TypeError
            pass

        self.key_info_table.setRowCount(0)  # 清空表格
        for key, value in key_info_params.items():
            row_position = self.key_info_table.rowCount()
            self.key_info_table.insertRow(row_position)
            self.key_info_table.setItem(row_position, 0, QTableWidgetItem(key))
            self.key_info_table.setItem(row_position, 1, QTableWidgetItem(str(value.get('x_num', 0.0))))
            self.key_info_table.setItem(row_position, 2, QTableWidgetItem(str(value.get('up_num', 0.0))))
            self.key_info_table.setItem(row_position, 3, QTableWidgetItem(str(value.get('do_num', 0.0))))

        # 重新连接信号
        self.key_info_table.itemChanged.connect(self.on_key_info_item_changed)

    def save_model_params(self):
        self.model_params = {
            'det_model': self.det_model.text(),
            'cls_model': self.cls_model.text(),
            'rec_model': self.rec_model.text(),
            'rec_label_file': self.rec_label_file.text(),
            'device': self.device.currentText(),
            'backend': self.backend.text(),
            'device_id': self.device_id.value(),
            'cpu_thread_num': self.cpu_thread_num.value(),
            'cls_bs': self.cls_bs.value(),
            'rec_bs': self.rec_bs.value(),
            'use_cls_model': self.use_cls_model.isChecked()
        }#'thread_num': self.thread_num.value(),'process_num': self.process_num.value()'use_multi_process': self.use_multi_process.isChecked()
        self.config['model_params'] = self.model_params
        self.save_config()
        print("模型参数已保存")
        
        # 重新初始化模型
        try:
            success = tl.initialize_model(self.model_params)
            if success:
                QMessageBox.information(self, "模型初始化", "模型已使用新参数重新初始化成功")
            else:
                QMessageBox.warning(self, "模型初始化", "模型重新初始化失败")
        except Exception as e:
            QMessageBox.critical(self, "模型初始化错误", f"模型重新初始化出错: {str(e)}")

    def load_model_params(self, params):
        if hasattr(self, 'det_model'):
            self.det_model.setText(params.get('det_model', ''))
        if hasattr(self, 'cls_model'):
            self.cls_model.setText(params.get('cls_model', ''))
        if hasattr(self, 'rec_model'):
            self.rec_model.setText(params.get('rec_model', ''))
        if hasattr(self, 'rec_label_file'):
            self.rec_label_file.setText(params.get('rec_label_file', ''))
        if hasattr(self, 'device'):
            self.device.setCurrentText(params.get('device', 'cpu'))
        if hasattr(self, 'backend'):
            self.backend.setText(params.get('backend', ''))
        if hasattr(self, 'device_id'):
            self.device_id.setValue(params.get('device_id', 0))
        if hasattr(self, 'cpu_thread_num'):
            self.cpu_thread_num.setValue(params.get('cpu_thread_num', 1))
        if hasattr(self, 'cls_bs'):
            self.cls_bs.setValue(params.get('cls_bs', 1))
        if hasattr(self, 'rec_bs'):
            self.rec_bs.setValue(params.get('rec_bs', 1))
        if hasattr(self, 'use_cls_model'):
            self.use_cls_model.setChecked(params.get('use_cls_model', False))
        # if hasattr(self, 'thread_num'):
        #     self.thread_num.setValue(params.get('thread_num', 1))
        # if hasattr(self, 'use_multi_process'):
        #     self.use_multi_process.setChecked(params.get('use_multi_process', False))
        # if hasattr(self, 'process_num'):
        #     self.process_num.setValue(params.get('process_num', 1))

    def get_model_params(self):
        if all(hasattr(self, attr) for attr in ['det_model', 'cls_model', 'rec_model', 'rec_label_file', 'device', 'backend', 'device_id', 'cpu_thread_num', 'cls_bs', 'rec_bs', 'use_cls_model']):#, 'thread_num', 'use_multi_process', 'process_num'
            return {
                'det_model': self.det_model.text(),
                'cls_model': self.cls_model.text(),
                'rec_model': self.rec_model.text(),
                'rec_label_file': self.rec_label_file.text(),
                'device': self.device.currentText(),
                'backend': self.backend.text(),
                'device_id': self.device_id.value(),
                'cpu_thread_num': self.cpu_thread_num.value(),
                'cls_bs': self.cls_bs.value(),
                'rec_bs': self.rec_bs.value(),
                'use_cls_model': self.use_cls_model.isChecked()
                # 'thread_num': self.thread_num.value(),
                # 'use_multi_process': self.use_multi_process.isChecked(),
                # 'process_num': self.process_num.value()
            }
        else:
            return {}

    def closeEvent(self, event):
        self.save_ui_state()
        super().closeEvent(event)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    main_window = MainWindow()
    main_window.show()
    sys.exit(app.exec_())