# mainwindow.py
import sys
import logging # 导入 logging 模块

from PySide6.QtWidgets import (QApplication, QMainWindow, QTreeWidgetItem,
                               QMessageBox, QHeaderView, QMenu, QInputDialog)
from PySide6.QtCore import Qt, QPoint

# 导入自动生成的 UI 文件
from ui_form import Ui_MainWindow
# 导入我们自己编写的系统信息管理模块
from system_info import SystemInfoManager
# 导入日志配置
from logger_config import setup_logging, logger # 导入日志配置函数和 logger 实例
# 导入磁盘操作模块
from disk_operations import DiskOperations

class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        # 设置日志输出到 QTextEdit
        setup_logging(self.ui.logOutputTextEdit)
        logger.info("应用程序启动。")

        # 初始化系统信息管理器和磁盘操作管理器
        self.system_manager = SystemInfoManager()
        self.disk_ops = DiskOperations()

        # 连接刷新按钮的信号到槽函数
        if hasattr(self.ui, 'refreshButton'):
            self.ui.refreshButton.clicked.connect(self.refresh_block_devices_info)
        else:
            logger.warning("Warning: refreshButton not found in UI. Please add it in form.ui and regenerate ui_form.py.")

        # 启用 treeWidget 的自定义上下文菜单
        self.ui.treeWidget_block_devices.setContextMenuPolicy(Qt.CustomContextMenu)
        self.ui.treeWidget_block_devices.customContextMenuRequested.connect(self.show_block_device_context_menu)

        # 初始化时刷新一次数据
        self.refresh_block_devices_info()
        logger.info("块设备信息已初始化加载。")

    def refresh_block_devices_info(self):
        """
        刷新块设备信息并显示在 QTreeWidget 中。
        """
        self.ui.treeWidget_block_devices.clear() # 清空现有内容

        # 定义所有要显示的列头和对应的 lsblk 字段名
        columns = [
            ("设备名", 'name'),
            ("类型", 'type'),
            ("大小", 'size'),
            ("挂载点", 'mountpoint'),
            ("文件系统", 'fstype'),
            ("只读", 'ro'),
            ("UUID", 'uuid'),
            ("PARTUUID", 'partuuid'),
            ("厂商", 'vendor'),
            ("型号", 'model'),
            ("序列号", 'serial'),
            ("主次号", 'maj:min'),
            ("父设备名", 'pkname'),
        ]

        headers = [col[0] for col in columns]
        self.field_keys = [col[1] for col in columns] # 保存字段键，供后续使用

        self.ui.treeWidget_block_devices.setColumnCount(len(headers)) # 确保列数正确
        self.ui.treeWidget_block_devices.setHeaderLabels(headers)

        # 调整列宽以适应内容
        for i in range(len(headers)):
            self.ui.treeWidget_block_devices.header().setSectionResizeMode(i, QHeaderView.ResizeToContents)

        try:
            devices = self.system_manager.get_block_devices()
            for dev in devices:
                self._add_device_to_tree(self.ui.treeWidget_block_devices, dev)

            # 自动调整列宽
            for i in range(len(headers)):
                self.ui.treeWidget_block_devices.resizeColumnToContents(i)
            logger.info("块设备信息刷新成功。")

        except Exception as e:
            QMessageBox.critical(self, "错误", f"刷新块设备信息失败: {e}")
            logger.error(f"刷新块设备信息失败: {e}")

    def _add_device_to_tree(self, parent_item, dev_data):
        """
        辅助函数，将单个设备及其子设备添加到 QTreeWidget。
        parent_item 可以是 QTreeWidget 本身，也可以是另一个 QTreeWidgetItem。
        """
        item = QTreeWidgetItem(parent_item)
        for i, key in enumerate(self.field_keys):
            value = dev_data.get(key)
            if key == 'ro': # 特殊处理布尔值
                item.setText(i, "是" if value else "否")
            elif value is None: # None 值显示为空字符串
                item.setText(i, "")
            else:
                item.setText(i, str(value))

        # 将原始设备数据存储在 item 的 data 属性中，方便后续操作时获取
        item.setData(0, Qt.UserRole, dev_data)

        # 如果有子设备（分区），也显示出来
        if 'children' in dev_data:
            for child in dev_data['children']:
                self._add_device_to_tree(item, child)
            item.setExpanded(True) # 默认展开父节点，以便看到分区

    def show_block_device_context_menu(self, pos: QPoint):
        """
        显示块设备列表的右键上下文菜单。
        """
        item = self.ui.treeWidget_block_devices.itemAt(pos)
        if item:
            dev_data = item.data(0, Qt.UserRole) # 获取存储的原始设备数据
            if not dev_data:
                logger.warning(f"无法获取设备 {item.text(0)} 的详细数据。")
                return

            device_name = dev_data.get('name')
            device_type = dev_data.get('type')
            mount_point = dev_data.get('mountpoint')

            menu = QMenu(self)

            # 挂载/卸载操作
            # 只有 'part' (分区) 和 'disk' (整个磁盘，但通常只挂载分区) 可以被挂载/卸载
            if device_type in ['part', 'disk']:
                if not mount_point or mount_point == '' or mount_point == 'N/A': # 未挂载
                    mount_action = menu.addAction(f"挂载 {device_name}...")
                    mount_action.triggered.connect(lambda: self._handle_mount(device_name))
                elif mount_point != '[SWAP]': # 已挂载且不是SWAP
                    unmount_action = menu.addAction(f"卸载 {device_name}")
                    unmount_action.triggered.connect(lambda: self._handle_unmount(device_name))

            # 分隔符，用于区分操作
            if menu.actions():
                menu.addSeparator()

            # 删除分区和格式化操作
            # 这些操作通常只针对 'part' (分区)
            if device_type == 'part':
                delete_action = menu.addAction(f"删除分区 {device_name}")
                delete_action.triggered.connect(lambda: self._handle_delete_partition(device_name))

                format_action = menu.addAction(f"格式化分区 {device_name}...")
                format_action.triggered.connect(lambda: self._handle_format_partition(device_name))

            if menu.actions(): # 只有当菜单中有动作时才显示
                menu.exec(self.ui.treeWidget_block_devices.mapToGlobal(pos))
            else:
                logger.info(f"设备 {device_name} 没有可用的操作。")
        else:
            logger.info("右键点击了空白区域。")

    def _handle_mount(self, device_name):
        """处理挂载操作，并刷新UI。"""
        if self.disk_ops.mount_partition(device_name):
            self.refresh_block_devices_info() # 操作成功后刷新UI

    def _handle_unmount(self, device_name):
        """处理卸载操作，并刷新UI。"""
        if self.disk_ops.unmount_partition(device_name):
            self.refresh_block_devices_info() # 操作成功后刷新UI

    def _handle_delete_partition(self, device_name):
        """处理删除分区操作，并刷新UI。"""
        if self.disk_ops.delete_partition(device_name):
            self.refresh_block_devices_info() # 操作成功后刷新UI

    def _handle_format_partition(self, device_name):
        """处理格式化分区操作，并刷新UI。"""
        if self.disk_ops.format_partition(device_name):
            self.refresh_block_devices_info() # 操作成功后刷新UI


if __name__ == "__main__":
    app = QApplication(sys.argv)
    widget = MainWindow()
    widget.show()
    sys.exit(app.exec())