#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UOS 全能系统升级助手
功能：
- 支持 upgrade / dist-upgrade 两种模式
- 预检查与备份提示
- 非交互式自动运行
- PyQt5 图形界面，实时显示输出
- needrestart 集成检测
"""

import sys
import os
import subprocess
import shutil
from datetime import datetime
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QTextEdit, QLabel,
                             QProgressBar, QMessageBox, QCheckBox, QRadioButton,
                             QButtonGroup, QGroupBox, QFileDialog)
from PyQt5.QtCore import QThread, pyqtSignal, Qt


# ==================== 工作线程：执行升级命令 ====================
class UpgradeThread(QThread):
    """后台执行升级命令，通过信号实时更新界面"""
    output_signal = pyqtSignal(str)      # 实时输出
    progress_signal = pyqtSignal(int)    # 进度更新
    finished_signal = pyqtSignal(bool, str)  # 完成信号(成功标志, 消息)

    def __init__(self, upgrade_mode, auto_confirm, backup_dir=None):
        super().__init__()
        self.upgrade_mode = upgrade_mode      # 'upgrade' or 'dist-upgrade'
        self.auto_confirm = auto_confirm      # True/False
        self.backup_dir = backup_dir          # 备份目录路径

    def run_command(self, cmd, description):
        """执行命令并实时输出（已在 root 环境下，无需 sudo）"""
        self.output_signal.emit(f"\n>>> {description}\n")
        self.output_signal.emit(f">>> 命令: {' '.join(cmd)}\n")

        # 非交互式设置环境变量
        env = os.environ.copy()
        if self.auto_confirm:
            env['DEBIAN_FRONTEND'] = 'noninteractive'
            # 添加 dpkg 选项，自动处理配置文件
            cmd = cmd + ['-o', 'Dpkg::Options::=--force-confdef',
                         '-o', 'Dpkg::Options::=--force-confold']

        try:
            process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                       stderr=subprocess.STDOUT,
                                       text=True, bufsize=1, env=env)
            for line in iter(process.stdout.readline, ''):
                if line:
                    self.output_signal.emit(line.rstrip())
            process.wait()
            if process.returncode == 0:
                self.output_signal.emit(f"\n✓ {description} 完成\n")
                return True
            else:
                self.output_signal.emit(f"\n✗ {description} 失败，返回码: {process.returncode}\n")
                return False
        except Exception as e:
            self.output_signal.emit(f"\n✗ 执行异常: {str(e)}\n")
            return False

    def run_needrestart(self):
        """运行 needrestart 检查需要重启的服务"""
        self.output_signal.emit("\n>>> 正在检查需要重启的服务...\n")
        try:
            # 检查哪些服务需要重启
            proc = subprocess.run(['needrestart', '-r', 'l'],
                                  capture_output=True, text=True)
            if proc.stdout:
                self.output_signal.emit(proc.stdout)
            if proc.stderr:
                self.output_signal.emit(proc.stderr)

            # 检查内核状态
            self.output_signal.emit("\n>>> 内核状态检查:\n")
            kernel = subprocess.run(['needrestart', '-b'],
                                    capture_output=True, text=True)
            if kernel.stdout:
                for line in kernel.stdout.splitlines():
                    if any(x in line for x in ['NEEDRESTART-KCUR', 'NEEDRESTART-KEXP', 'NEEDRESTART-SVC']):
                        self.output_signal.emit(f"  {line}\n")
            return True
        except FileNotFoundError:
            self.output_signal.emit("[提示] needrestart 未安装，跳过服务检查。可用 'sudo apt install needrestart' 安装\n")
            return False

    def create_backup(self):
        """创建软件源备份（/etc/apt/sources.list 和 sources.list.d）"""
        if not self.backup_dir:
            return True

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_path = os.path.join(self.backup_dir, f"sources_backup_{timestamp}")

        self.output_signal.emit(f"\n>>> 正在备份软件源配置到: {backup_path}\n")
        try:
            os.makedirs(backup_path, exist_ok=True)
            # 备份 sources.list.d 目录
            shutil.copytree('/etc/apt/sources.list.d',
                            os.path.join(backup_path, 'sources.list.d'),
                            dirs_exist_ok=True)   # Python 3.8+ 允许覆盖
            # 备份 sources.list 文件
            shutil.copy2('/etc/apt/sources.list',
                         os.path.join(backup_path, 'sources.list'))
            self.output_signal.emit("✓ 软件源备份完成\n")
            return True
        except Exception as e:
            self.output_signal.emit(f"✗ 备份失败: {str(e)}\n")
            return False

    def run(self):
        """线程主函数"""
        success = True

        # 0. 创建备份（如果需要）
        if self.backup_dir:
            self.create_backup()

        # 1. 更新软件源
        self.progress_signal.emit(10)
        if not self.run_command(['apt', 'update'], '更新软件源列表'):
            self.finished_signal.emit(False, "软件源更新失败，请检查网络或源配置")
            return

        # 2. 预检查：显示可升级的软件包
        self.progress_signal.emit(30)
        self.output_signal.emit("\n>>> 检查可升级的软件包...\n")
        try:
            result = subprocess.run(['apt', 'list', '--upgradable'],
                                    capture_output=True, text=True)
            upgradable_lines = [l for l in result.stdout.splitlines()
                                if l and 'Listing...' not in l]
            if upgradable_lines:
                self.output_signal.emit(f"发现 {len(upgradable_lines)} 个软件包可升级：\n")
                for line in upgradable_lines[:10]:
                    self.output_signal.emit(f"  {line}\n")
                if len(upgradable_lines) > 10:
                    self.output_signal.emit(f"  ... 还有 {len(upgradable_lines)-10} 个包\n")
            else:
                self.output_signal.emit("当前系统已是最新，无需升级。\n")
        except Exception as e:
            self.output_signal.emit(f"无法获取升级列表: {e}\n")

        # 3. 执行系统升级
        self.progress_signal.emit(50)
        upgrade_cmd = ['apt', self.upgrade_mode, '-y']
        upgrade_desc = '执行系统升级' if self.upgrade_mode == 'upgrade' else '执行完整系统升级(dist-upgrade)'

        if not self.run_command(upgrade_cmd, upgrade_desc):
            self.finished_signal.emit(False, f"{upgrade_desc}失败")
            return

        # 4. 清理无用依赖
        self.progress_signal.emit(80)
        self.run_command(['apt', 'autoremove', '-y'], '清理无用依赖包')
        self.run_command(['apt', 'autoclean'], '清理软件包缓存')

        # 5. needrestart 检查
        self.progress_signal.emit(90)
        self.run_needrestart()

        # 6. 完成
        self.progress_signal.emit(100)
        self.finished_signal.emit(True, "系统升级完成！")
        self.output_signal.emit("\n========== 升级流程全部完成 ==========\n")


# ==================== 主窗口 ====================
class UpgradeWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.upgrade_thread = None
        self.initUI()

    def initUI(self):
        self.setWindowTitle("UOS 全能系统升级助手")
        self.setMinimumSize(800, 600)

        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # ===== 模式选择区 =====
        mode_group = QGroupBox("升级模式")
        mode_layout = QHBoxLayout()
        self.radio_upgrade = QRadioButton("标准升级 (upgrade)")
        self.radio_dist = QRadioButton("完整升级 (dist-upgrade) *推荐")
        self.radio_upgrade.setChecked(True)
        self.radio_dist.setToolTip("可以处理内核更换和复杂的依赖关系变更")
        mode_layout.addWidget(self.radio_upgrade)
        mode_layout.addWidget(self.radio_dist)
        mode_group.setLayout(mode_layout)
        main_layout.addWidget(mode_group)

        # 分组（便于扩展）
        self.mode_group = QButtonGroup()
        self.mode_group.addButton(self.radio_upgrade, 0)
        self.mode_group.addButton(self.radio_dist, 1)

        # ===== 预检查与备份设置 =====
        backup_group = QGroupBox("预检查与备份")
        backup_layout = QVBoxLayout()

        self.check_before_btn = QCheckBox("升级前备份软件源配置 (推荐)")
        self.check_before_btn.setChecked(True)
        backup_layout.addWidget(self.check_before_btn)

        # 备份路径选择
        path_layout = QHBoxLayout()
        path_layout.addWidget(QLabel("备份目录:"))
        self.backup_path_edit = QTextEdit()
        self.backup_path_edit.setMaximumHeight(30)
        self.backup_path_edit.setText(os.path.expanduser("~/uos_upgrade_backup"))
        path_layout.addWidget(self.backup_path_edit)

        self.browse_btn = QPushButton("浏览...")
        self.browse_btn.clicked.connect(self.browse_backup_dir)
        path_layout.addWidget(self.browse_btn)
        backup_layout.addLayout(path_layout)

        backup_group.setLayout(backup_layout)
        main_layout.addWidget(backup_group)

        # ===== 非交互式设置 =====
        auto_group = QGroupBox("自动化设置")
        auto_layout = QHBoxLayout()
        self.auto_confirm_btn = QCheckBox("非交互式运行 (自动确认所有提示)")
        self.auto_confirm_btn.setChecked(True)
        self.auto_confirm_btn.setToolTip("启用 DEBIAN_FRONTEND=noninteractive，跳过交互式配置")
        auto_layout.addWidget(self.auto_confirm_btn)
        auto_group.setLayout(auto_layout)
        main_layout.addWidget(auto_group)

        # ===== 控制按钮 =====
        btn_layout = QHBoxLayout()
        self.start_btn = QPushButton("开始升级")
        self.start_btn.clicked.connect(self.start_upgrade)
        self.start_btn.setStyleSheet("background-color: #4CAF50; font-weight: bold; padding: 8px;")

        self.stop_btn = QPushButton("停止")
        self.stop_btn.clicked.connect(self.stop_upgrade)
        self.stop_btn.setEnabled(False)
        self.stop_btn.setStyleSheet("background-color: #f44336; padding: 8px;")

        btn_layout.addWidget(self.start_btn)
        btn_layout.addWidget(self.stop_btn)
        main_layout.addLayout(btn_layout)

        # ===== 进度条 =====
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        main_layout.addWidget(self.progress_bar)

        # ===== 输出区域 =====
        output_label = QLabel("执行日志 (实时输出):")
        main_layout.addWidget(output_label)

        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        self.output_text.setFontFamily("Monospace")
        main_layout.addWidget(self.output_text)

        # 窗口底部状态栏
        self.statusBar().showMessage("就绪 - 点击“开始升级”")

    def browse_backup_dir(self):
        """浏览备份目录"""
        dir_path = QFileDialog.getExistingDirectory(self, "选择备份目录",
                                                    self.backup_path_edit.toPlainText())
        if dir_path:
            self.backup_path_edit.setText(dir_path)

    def log(self, message):
        """在输出区域追加消息并自动滚动到底部"""
        self.output_text.append(message)
        cursor = self.output_text.textCursor()
        cursor.movePosition(cursor.End)
        self.output_text.setTextCursor(cursor)

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

    def start_upgrade(self):
        """开始升级流程"""
        # 获取用户配置
        upgrade_mode = 'dist-upgrade' if self.radio_dist.isChecked() else 'upgrade'
        auto_confirm = self.auto_confirm_btn.isChecked()
        backup_dir = self.backup_path_edit.toPlainText().strip() if self.check_before_btn.isChecked() else None

        # 确认提示
        msg = f"即将执行{'完整' if upgrade_mode == 'dist-upgrade' else '标准'}系统升级。\n"
        if upgrade_mode == 'dist-upgrade':
            msg += "⚠️ 注意: dist-upgrade 可能会安装新包或卸载有冲突的包。\n"
        if auto_confirm:
            msg += "🤖 非交互模式已启用，过程不会出现交互提示。\n"
        msg += "\n是否继续？"

        reply = QMessageBox.question(self, "确认升级", msg,
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)
        if reply != QMessageBox.Yes:
            return

        # 清空输出区
        self.output_text.clear()
        self.progress_bar.setValue(0)

        # 创建并启动线程
        self.upgrade_thread = UpgradeThread(upgrade_mode, auto_confirm, backup_dir)
        self.upgrade_thread.output_signal.connect(self.log)
        self.upgrade_thread.progress_signal.connect(self.update_progress)
        self.upgrade_thread.finished_signal.connect(self.on_upgrade_finished)

        self.upgrade_thread.start()

        # 更新界面状态
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.statusBar().showMessage("升级进行中...")

    def stop_upgrade(self):
        """停止升级（终止线程）"""
        if self.upgrade_thread and self.upgrade_thread.isRunning():
            reply = QMessageBox.question(self, "确认停止",
                                         "正在运行的升级可能被中断，是否确定停止？",
                                         QMessageBox.Yes | QMessageBox.No,
                                         QMessageBox.No)
            if reply == QMessageBox.Yes:
                self.upgrade_thread.terminate()
                self.upgrade_thread.wait()
                self.log("\n⚠️ 用户手动停止了升级流程。")
                self.on_upgrade_finished(False, "升级已被用户中断")

    def on_upgrade_finished(self, success, message):
        """升级完成回调"""
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.progress_bar.setValue(100 if success else 0)

        if success:
            self.statusBar().showMessage("升级完成")
            QMessageBox.information(self, "完成", message)
        else:
            self.statusBar().showMessage("升级失败")
            QMessageBox.critical(self, "错误", f"{message}\n请查看日志信息。")

        # 如果升级成功且选择了完整升级，提示可能需要的重启
        if success and self.radio_dist.isChecked():
            QMessageBox.information(self, "提示",
                                    "完整升级已完成。\n"
                                    "如系统运行异常，建议重启系统以确保所有更新生效。")


# ==================== 程序入口 ====================
def main():
    app = QApplication(sys.argv)
    window = UpgradeWindow()
    window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()