#!/usr/bin/env python3
"""
DCU Performance Analyzer Tool
A comprehensive performance analysis tool for DCU environment checking and monitoring.
"""

import os
import sys
import json
import time
import logging
import argparse
import subprocess
import platform
import tarfile
import shutil
import re
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed
from collections import defaultdict

__version__ = "1.0.0"
__author__ = "DCU Performance Team"


def analyze_region_line(region_line: str):
    match = re.search(r"Memory at (\S+)", region_line)
    if not match:
        return "ERROR", "无法解析BAR地址", False
    address = match.group(1)

    if address == "unassigned":
        return "FAIL", "BAR未分配，可能MMIO不足", True

    if address == "<ignored>":
        return "FAIL", "BAR被忽略，需开启Above 4G或pcie=realloc", True

    try:
        addr_int = int(address, 16)
        if addr_int > 0xFFFFFFFFFFF:
            return "WARNING", "BAR超过44bit", False
    except:
        pass

    return "PASS", "正常", False

def get_cpu_numa_topology():
    topology = {}
    try:
        for node in os.listdir("/sys/devices/system/node/"):
            if node.startswith("node"):
                node_id = int(node.replace("node", ""))
                cpulist_path = f"/sys/devices/system/node/{node}/cpulist"
                if os.path.exists(cpulist_path):
                    cpulist = open(cpulist_path).read().strip()
                    topology[node_id] = cpulist
    except:
        pass
    return topology

@dataclass
class CheckResult:
    """检查结果数据结构"""
    module: str
    status: str  # 'PASS', 'FAIL', 'WARNING', 'INFO'
    message: str
    details: Dict
    timestamp: str
    execution_time: float

class DCUPerformanceAnalyzer:
    """DCU性能分析工具主类"""
    
    def __init__(self, output_dir: str = None, debug: bool = False, quiet: bool = True):
        self.start_time = datetime.now()
        self.output_dir = output_dir or f"dcu_analysis_{self.start_time.strftime('%Y%m%d_%H%M%S')}"
        self.debug = debug
        self.quiet = quiet
        self.results: List[CheckResult] = []
        
        os.makedirs(self.output_dir, exist_ok=True)
        os.makedirs(f"{self.output_dir}/logs", exist_ok=True)
        os.makedirs(f"{self.output_dir}/reports", exist_ok=True)
        os.makedirs(f"{self.output_dir}/data", exist_ok=True)

        self.logger = self._setup_logging()
        # self.scripts_path = Path("/data1/csy/project/small_model/dcuprofiletools/env_check/dcu_env_check-main")
        
    def _setup_logging(self) -> logging.Logger:
        logger = logging.getLogger('DCUAnalyzer')
        logger.setLevel(logging.DEBUG if self.debug else logging.INFO)
        log_dir = f'{self.output_dir}/logs'
        os.makedirs(log_dir, exist_ok=True)
        fh = logging.FileHandler(f'{log_dir}/analysis.log')
        fh.setLevel(logging.DEBUG)
        if self.quiet:
            ch = logging.StreamHandler()
            ch.setLevel(logging.INFO)
            logger.addHandler(ch)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        logger.addHandler(fh)
        return logger
    
    def run_command(self, cmd: str, timeout: int = 300, cwd: Optional[str] = None) -> Tuple[int, str, str]:
        effective_cwd = cwd
        self.logger.debug(f"执行命令: {cmd} in {effective_cwd}")
        try:
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout, cwd=effective_cwd)
            return result.returncode, result.stdout, result.stderr
        except subprocess.TimeoutExpired:
            self.logger.error(f"命令超时: {cmd}")
            return -1, "", "Command timeout"
        except Exception as e:
            self.logger.error(f"命令执行失败: {cmd}, 错误: {str(e)}")
            return -1, "", str(e)

    def check_system_info(self) -> CheckResult:
        start_time = time.time()
        self.logger.info("开始系统信息检查...")
        try:
            commands = {
                "cpu_info.txt": "lscpu",
                "mem_info.txt": "free -h && echo '\n' && dmidecode -t memory",
                "network.txt": "ip a && echo '\n' && lspci -nn | grep -i -E \"eth|mellanox\"",
                "os_info.txt": [
                    "cat /etc/os-release", "uname -a", "cat /proc/cmdline",
                    "numactl -H", "rpm -qf $(which ldd)", "ldd --version",
                    "strings $(find /usr/ -name libc.so.6) | grep ^GLIBC_",
                    "strings $(find /usr -name libstdc++.so.6) | grep GLIBCXX"
                ],
            }
            details = {}
            for filename, command in commands.items():
                full_path = f"{self.output_dir}/data/{filename}"
                if isinstance(command, list):
                    with open(full_path, "w") as f:
                        for cmd_item in command:
                            ret_code, stdout, stderr = self.run_command(cmd_item)
                            f.write(f"### {cmd_item} ###\n")
                            f.write(stdout if ret_code == 0 else f"ERROR: {stderr}")
                            f.write("\n\n")
                else:
                    ret_code, stdout, stderr = self.run_command(command)
                    with open(full_path, "w") as f:
                        f.write(stdout if ret_code == 0 else f"ERROR: {stderr}")
                details[filename] = "saved"
            return CheckResult("System Information", "PASS", "系统信息收集完成", details, datetime.now().isoformat(), time.time() - start_time)
        except Exception as e:
            return CheckResult("System Information", "FAIL", f"系统信息收集失败: {str(e)}", {"error": str(e)}, datetime.now().isoformat(), time.time() - start_time)

    def check_pcie_devices(self) -> CheckResult:
        start_time = time.time()
        self.logger.info("开始PCIe设备检查...")

        log_dir = f"{self.output_dir}/data/pcie_info"
        os.makedirs(log_dir, exist_ok=True)

        summary_log = f"{self.output_dir}/data/pcie_dcu.log"
        os.makedirs(f"{self.output_dir}/data", exist_ok=True)

        details = {
            "total_dcu": 0,
            "models": {},
            "devices": {},
            "bar_errors": 0,
            "mmio_issue_count": 0,
            "link_downgrade": 0,
            "cross_numa_count": 0,
            "pcie_topology": ""
        }

        dcu_id_map = {
            "54b7": ["Z100"],
            "55b7": ["Z100L"],
            "62b7": ["K100"],
            "6210": ["K100-AI"],
            "6211": ["K100-AI-ECO"],
            "6320": ["BW1000"]
        }

        ret, lspci_full, _ = self.run_command("lspci -D -nn")
        with open(f"{log_dir}/pcie_info.log", "w") as f:
            f.write(lspci_full)

        ret, lspci_tree, _ = self.run_command("lspci -t")
        details["pcie_topology"] = lspci_tree
        with open(f"{log_dir}/pcie_vt.log", "w") as f:
            f.write(lspci_tree)

        ret, lspci_all_verbose, _ = self.run_command("lspci -vvv")
        with open(f"{log_dir}/pcie_more.log", "w") as f:
            f.write(lspci_all_verbose)

        cpu_topology = get_cpu_numa_topology()
        try:
            with open(summary_log, "w", encoding="utf-8") as logf:

                logf.write("===== PCIe 拓扑 =====\n")
                logf.write(lspci_tree + "\n")
                for line in lspci_full.splitlines():
                    if not re.search(r"display|co-processor", line, re.IGNORECASE):
                        continue

                    match = re.search(r"\[(\w+):(\w+)\]", line)
                    if not match:
                        continue

                    device_id = match.group(2).lower()
                    if device_id not in dcu_id_map:
                        continue

                    bdf = line.split()[0]
                    model_name = ",".join(dcu_id_map[device_id])

                    details["total_dcu"] += 1
                    details["models"].setdefault(model_name, 0)
                    details["models"][model_name] += 1

                    # 获取单设备verbose
                    ret, verbose, _ = self.run_command(f"lspci -vv -s {bdf}")

                    # 单设备日志单独保存
                    with open(f"{log_dir}/lspci_{bdf.replace(':','_').replace('.','_')}.log", "w") as f:
                        f.write(verbose)

                    # 链路
                    cur_speed = re.search(r"LnkSta:.*Speed (\S+)", verbose)
                    cur_width = re.search(r"LnkSta:.*Width x(\d+)", verbose)
                    max_speed = re.search(r"LnkCap:.*Speed (\S+)", verbose)
                    max_width = re.search(r"LnkCap:.*Width x(\d+)", verbose)

                    downgrade = False
                    if cur_speed and max_speed and cur_speed.group(1) != max_speed.group(1):
                        downgrade = True
                    if cur_width and max_width and cur_width.group(1) != max_width.group(1):
                        downgrade = True
                    if downgrade:
                        details["link_downgrade"] += 1

                    # BAR
                    bar_regions = re.findall(r"Region \d+:.*", verbose)
                    bar_analysis = []

                    for region in bar_regions:
                        status, msg, mmio_issue = analyze_region_line(region)
                        if status == "FAIL":
                            details["bar_errors"] += 1
                        if mmio_issue:
                            details["mmio_issue_count"] += 1

                        bar_analysis.append({
                            "region": region.strip(),
                            "status": status,
                            "message": msg
                        })

                    # NUMA
                    numa_path = f"/sys/bus/pci/devices/{bdf}/numa_node"
                    numa_node = "unknown"
                    if os.path.exists(numa_path):
                        numa_node = open(numa_path).read().strip()

                    cross_numa = False
                    if numa_node != "unknown":
                        node_id = int(numa_node)
                        if node_id not in cpu_topology:
                            cross_numa = True
                            details["cross_numa_count"] += 1

                    device_info = {
                        "device_id": device_id,
                        "model": model_name,
                        "numa_node": numa_node,
                        "cross_numa": cross_numa,
                        "downgrade": downgrade,
                        "bar_regions": bar_analysis
                    }

                    details["devices"][bdf] = device_info

                    logf.write("\n" + "=" * 60 + "\n")
                    logf.write(f"DCU设备: {bdf}\n")
                    for k, v in device_info.items():
                        logf.write(f"{k}: {v}\n")

            # 最终判定
            status = "PASS"

            if details["mmio_issue_count"] > 0:
                status = "FAIL"
            elif details["bar_errors"] > 0:
                status = "FAIL"
            elif details["cross_numa_count"] > 0:
                status = "WARNING"
            elif details["link_downgrade"] > 0:
                status = "WARNING"

            return CheckResult(
                module="PCIe Devices",
                status=status,
                message=f"检测到 {details['total_dcu']} 张 DCU",
                details=details,
                timestamp=datetime.now().isoformat(),
                execution_time=time.time() - start_time
            )

        except Exception as e:
            return CheckResult(
                module="PCIe Devices",
                status="FAIL",
                message=str(e),
                details={"error": str(e)},
                timestamp=datetime.now().isoformat(),
                execution_time=time.time() - start_time
            )
            
    def check_driver_status(self) -> CheckResult:

        log_dir = f"{self.output_dir}/data/"
        os.makedirs(log_dir, exist_ok=True)

        start_time = time.time()
        self.logger.info("开始驱动状态检查...")
        
        details = {
            "driver_loaded": False,
            "hydcu_driver_loaded": False,
            "iommu_status": "",
            "vfio_attached": False,
            "missing_kos": [],
            "driver_info": {}
        }
        
        messages = []
        
        ret_code, stdout, stderr = self.run_command("lsmod | grep hycu")
        if ret_code == 0 and stdout.strip():
            details["hydcu_driver_loaded"] = True
            details["driver_loaded"] = True
            messages.append("hydcu 驱动已成功加载")
            messages.append(f"lsmod 输出: {stdout.strip()}")
        else:
            messages.append("hydcu 驱动未加载")
        
        if details["driver_loaded"]:
            ret_code, stdout, stderr = self.run_command("/opt/hyhal/bin/hy-smi -a")

            with open(f"{log_dir}/hy_smi_info.txt", "w") as f:
                f.write(stdout)

            if ret_code == 0 and stdout.strip():
                driver_match = re.search(r"Driver Version:\s*([^\s]+)", stdout)
                vbios_matches = re.findall(r"VBIOS version:\s*([^\s]+)", stdout)
                driver_version = driver_match.group(1) if driver_match else "Not Found"
                details["driver_info"]["version"] = driver_version
                details["driver_info"]["vbios"] = vbios_matches

                messages.append("成功获取 hydcu 驱动信息")
            else:
                messages.append("无法获取 hydcu 驱动信息")
                if stderr:
                    messages.append(f"错误: {stderr}")
        
        ret_code, stdout, stderr = self.run_command("lsmod | grep iommu_v2")
        if ret_code == 0 and stdout.strip():
            details["iommu_status"] = "iommu_v2_loaded"
            messages.append("IOMMU v2 驱动已加载")
        else:
            ret_code, stdout, stderr = self.run_command("lsmod | grep amd_iommu_v2")
            if ret_code == 0 and stdout.strip():
                details["iommu_status"] = "amd_iommu_v2_loaded"
                messages.append("AMD IOMMU v2 驱动已加载")
            else:
                # 检查驱动是否存在但未加载
                ret_code, stdout, stderr = self.run_command("modinfo iommu_v2")
                if ret_code == 0 and stdout.strip():
                    details["iommu_status"] = "iommu_v2_present_not_loaded"
                    messages.append("IOMMU v2 驱动存在但未加载")
                else:
                    ret_code, stdout, stderr = self.run_command("modinfo amd_iommu_v2")
                    if ret_code == 0 and stdout.strip():
                        details["iommu_status"] = "amd_iommu_v2_present_not_loaded"
                        messages.append("AMD IOMMU v2 驱动存在但未加载")
                    else:
                        details["iommu_status"] = "no_iommu_driver"
                        messages.append("系统中未找到 IOMMU 驱动")
        
        ret_code, stdout, stderr = self.run_command("lsmod | grep vfio_pci")
        if ret_code == 0 and stdout.strip():
            details["vfio_attached"] = True
            messages.append("检测到 VFIO-PCI 驱动，某些设备可能已附加到虚拟机")
        else:
            ret_code, stdout, stderr = self.run_command("lsmod | grep vfio-pci")
            if ret_code == 0 and stdout.strip():
                details["vfio_attached"] = True
                messages.append("检测到 VFIO-PCI 驱动，某些设备可能已附加到虚拟机")
        
        ko_list = [
            "hycu.ko",
            "hycu-sched.ko", 
            "hydrm_ttm_helper.ko",
            "hy-extra.ko",
            "hykcl.ko",
            "hyttm.ko"
        ]
        ko_dir = "/opt/hyhal/dkms"
        
        for ko in ko_list:
            ko_path = os.path.join(ko_dir, ko)
            if not os.path.isfile(ko_path):
                details["missing_kos"].append(ko)
                messages.append(f"缺失驱动文件: {ko}")
            elif not os.access(ko_path, os.R_OK):
                details["missing_kos"].append(ko)
                messages.append(f"驱动文件无读取权限: {ko}")
        
        status = "PASS"
        
        if details["missing_kos"]:
            status = "FAIL"
        elif not details["driver_loaded"]:
            status = "FAIL"
        elif details["iommu_status"] in ["no_iommu_driver", "iommu_v2_present_not_loaded", "amd_iommu_v2_present_not_loaded"]:
            status = "WARNING"
        elif details["vfio_attached"]:
            status = "WARNING"
        
        return CheckResult(
            module="Driver Status",
            status=status,
            message="驱动检查完成",
            details={
                "summary": "\n".join(messages),
                **details
            },
            timestamp=datetime.now().isoformat(),
            execution_time=time.time() - start_time
        )

    def check_system_logs(self, log_age_hours: int = 24, log_size_limit_mb: int = 10) -> CheckResult:
        start_time = time.time()
        self.logger.info("开始系统日志检查...")
        try:
            details = {}
            ret_code, dmesg_output, _ = self.run_command("dmesg -T")
            if ret_code == 0: 
                with open(f"{self.output_dir}/data/dmesg.log", "w") as f: f.write(dmesg_output)
                details["dmesg.log"] = "saved"
            syslog_path = "/var/log/syslog" if os.path.exists("/var/log/syslog") else "/var/log/messages" if os.path.exists("/var/log/messages") else None
            dest_path = f"{self.output_dir}/data/system.log"
            if syslog_path:
                try:
                    if os.path.getsize(syslog_path) / (1024*1024) > log_size_limit_mb:
                        with open(dest_path, "w") as f: f.write("[日志文件超过大小限制未采集]")
                        details["system.log"] = "skipped_too_large"
                    else: 
                        shutil.copy(syslog_path, dest_path)
                        details["system.log"] = "copied"
                except Exception as e:
                    with open(dest_path, "w") as f: f.write(f"无权限读取日志: {e}")
                    details["system.log"] = "permission_denied"
            else:
                ret_code, journal_logs, _ = self.run_command(f"journalctl --since '{log_age_hours} hours ago'")
                with open(dest_path, "w") as f: f.write(journal_logs if ret_code == 0 else "无法获取系统日志")
                details["system.log"] = "journal_saved" if ret_code == 0 else "journal_failed"
            ret_code, sme_log, _ = self.run_command(f"grep -i sme {self.output_dir}/data/dmesg.log")
            # with open(f"{self.output_dir}/data/sme.log", "w") as f: f.write(sme_log)
            details["sme_status"] = "非CSV场景，需要BIOS关闭SME设置" if sme_log.strip() else "OS SME目前没有打开, 非CSV场景下，该状态正常"
            return CheckResult("System Logs", "PASS", "系统日志收集完成", details, datetime.now().isoformat(), time.time() - start_time)
        except Exception as e:
            return CheckResult("System Logs", "FAIL", f"系统日志检查失败: {str(e)}", {"error": str(e)}, datetime.now().isoformat(), time.time() - start_time)

    def check_hardware_info(self) -> CheckResult:
        start_time = time.time()
        self.logger.info("开始硬件信息检查...")
        try:
            commands = {"hardware.txt": ["ipmitool fru", "ipmitool mc info", "dmidecode -s system-product-name", "dmidecode -t bios"]}
            details = {}
            for filename, command_list in commands.items():
                full_path = f"{self.output_dir}/data/{filename}"
                with open(full_path, "w") as f:
                    for cmd_item in command_list:
                        ret_code, stdout, stderr = self.run_command(cmd_item)
                        f.write(f"### {cmd_item} ###\n")
                        f.write(stdout if ret_code == 0 else f"ERROR: {stderr}")
                        f.write("\n\n")
                details[filename] = "saved"
            ret_code, serial_num, _ = self.run_command("dmidecode -t 2 | grep -i 'Serial Number' | awk '{print $3}'")
            if ret_code == 0 and serial_num.strip():
                details["serial_number"] = serial_num.strip()
                if "AS" in serial_num.strip(): details["board_warning"] = "检测到主板型号过旧，满负载情况可能出现掉卡"
                elif "BH" in serial_num.strip(): details["board_status"] = "主板型号符合要求"
            return CheckResult("Hardware Information", "PASS", "硬件信息收集完成", details, datetime.now().isoformat(), time.time() - start_time)
        except Exception as e:
            return CheckResult("Hardware Information", "FAIL", f"硬件信息检查失败: {str(e)}", {"error": str(e)}, datetime.now().isoformat(), time.time() - start_time)

    def check_pkg_status(self, auto_install: bool = False) -> CheckResult:
        start_time = time.time()
        self.logger.info("开始软件包状态检查...")
        
        details = {
            "missing_packages": [],
            "installed_packages": [],
            "package_manager": None,
            "auto_install_attempted": auto_install,
            "auto_install_success": False
        }
        
        messages = []
        
        # 检测系统类型和包管理器
        if os.path.exists("/etc/debian_version"):
            details["package_manager"] = "apt"
            pkgs_info = {
                "dmidecode": "dmidecode",
                "lshw": "lshw", 
                "lspci": "pciutils",
                "numactl": "numactl-devel",
                "ipmitool": "ipmitool",
                "locate": "locate"
            }
        elif os.path.exists("/etc/redhat-release") or os.path.exists("/etc/centos-release"):
            details["package_manager"] = "yum"
            pkgs_info = {
                "dmidecode": "dmidecode",
                "lshw": "lshw",
                "lspci": "pciutils", 
                "numactl": "numactl-dev",
                "ipmitool": "ipmitool",
                "locate": "mlocate"
            }
        else:
            details["package_manager"] = "unknown"
            pkgs_info = {
                "dmidecode": "dmidecode",
                "lshw": "lshw",
                "lspci": "pciutils",
                "numactl": "numactl-devel",
                "ipmitool": "ipmitool",
                "locate": "locate"
            }
        
        # 检查每个命令是否可用
        commands_to_check = ["dmidecode", "lshw", "lspci", "numactl", "ipmitool", "locate"]
        
        for cmd in commands_to_check:
            ret_code, stdout, stderr = self.run_command(f"command -v {cmd}")
            if ret_code == 0 and stdout.strip():
                details["installed_packages"].append(cmd)
                messages.append(f"✓ {cmd} 已安装")
            else:
                details["missing_packages"].append(cmd)
                pkg_name = pkgs_info.get(cmd, cmd)
                messages.append(f"✗ {cmd} 未安装 (包名: {pkg_name})")
        
        # 状态判定
        status = "PASS" if not details["missing_packages"] else "WARNING"
        
        # 如果需要自动安装且缺失包存在
        if auto_install and details["missing_packages"] and details["package_manager"] in ["apt", "yum"]:
            messages.append("尝试自动安装缺失的软件包...")
            
            if details["package_manager"] == "apt":
                install_cmd = f"sudo apt-get update && sudo apt-get install -y {' '.join([pkgs_info[pkg] for pkg in details['missing_packages']])}"
            else:  # yum
                install_cmd = f"sudo yum install -y {' '.join([pkgs_info[pkg] for pkg in details['missing_packages']])}"
            
            ret_code, stdout, stderr = self.run_command(install_cmd)
            if ret_code == 0:
                details["auto_install_success"] = True
                messages.append("自动安装成功")
                
                # 重新检查
                still_missing = []
                for cmd in details["missing_packages"]:
                    ret_code, stdout, stderr = self.run_command(f"command -v {cmd}")
                    if ret_code != 0:
                        still_missing.append(cmd)
                
                if not still_missing:
                    status = "PASS"
                    messages.append("所有软件包安装完成")
                else:
                    messages.append(f"仍有软件包安装失败: {', '.join(still_missing)}")
            else:
                messages.append(f"自动安装失败: {stderr}")
        
        return CheckResult(
            module="Package Status",
            status=status,
            message="软件包检查完成",
            details={
                "summary": "\n".join(messages),
                **details
            },
            timestamp=datetime.now().isoformat(),
            execution_time=time.time() - start_time
        )

    # def check_rccl_info(self) -> CheckResult:
    #     start_time = time.time()
    #     self.logger.info("开始RCCL信息检查...")

    #     log_dir = f"{self.output_dir}/data/"
    #     os.makedirs(log_dir, exist_ok=True)


    def run_all_checks(self, log_age: int = 24, log_size: int = 10, auto_install_pkg: bool = False) -> List[CheckResult]:
        self.logger.info("开始执行所有检查模块...")
        checks = [
            self.check_system_info, self.check_pcie_devices, self.check_driver_status,
            lambda: self.check_system_logs(log_age, log_size),
            self.check_hardware_info,
            lambda: self.check_pkg_status(auto_install_pkg)
        ]
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = {executor.submit(check): check.__name__ for check in checks}
            for future in as_completed(futures):
                check_name = futures[future]
                try:
                    result = future.result()
                    self.results.append(result)
                    self.logger.info(f"{check_name} 完成 - 状态: {result.status}")
                except Exception as e:
                    self.logger.error(f"{check_name} 执行失败: {str(e)}")
        return self.results

    def generate_report(self) -> str:
        self.logger.info("生成分析报告...")
        report_file = f"{self.output_dir}/reports/analysis_report.json"
        summary_file = f"{self.output_dir}/reports/analysis_summary.txt"
        status_counts = {"PASS": 0, "FAIL": 0, "WARNING": 0, "INFO": 0}
        total_execution_time = sum(r.execution_time for r in self.results)
        for result in self.results: status_counts[result.status] += 1
        report = {
            "metadata": {"tool_version": __version__, "analysis_date": self.start_time.isoformat(), "hostname": platform.node(), "total_execution_time": total_execution_time, "output_directory": self.output_dir},
            "summary": {"total_checks": len(self.results), "status_distribution": status_counts, "overall_status": "PASS" if status_counts["FAIL"] == 0 else "FAIL"},
            "results": [r.__dict__ for r in self.results]
        }
        with open(report_file, "w", encoding="utf-8") as f: json.dump(report, f, indent=2, ensure_ascii=False)
        with open(summary_file, "w", encoding="utf-8") as f:
            f.write(f'DCU性能分析报告摘要\n分析时间: {self.start_time.strftime("%Y-%m-%d %H:%M:%S")}\n')
            f.write(f"整体状态: {report['summary']['overall_status']}\n")
            for result in self.results:
                f.write(f"\n模块: {result.module} - 状态: {result.status}\n消息: {result.message}\n")
        self.logger.info(f"报告生成完成: {report_file}")
        return report_file

    def create_package(self) -> str:
        self.logger.info("创建分析结果包...")
        package_name = f"{self.output_dir}.tar.gz"
        with tarfile.open(package_name, "w:gz") as tar: tar.add(self.output_dir, arcname=os.path.basename(self.output_dir))
        self.logger.info(f"分析结果包创建完成: {package_name}")
        return package_name

    def run_analysis(self, checks: List[str] = None, log_age: int = 24, log_size: int = 10, auto_install_pkg: bool = False) -> str:
        self.logger.info("开始DCU性能分析...")
        try:
            available_checks = {
                "pkg": lambda: self.check_pkg_status(auto_install_pkg),
                "system": self.check_system_info, "pcie": self.check_pcie_devices, "driver": self.check_driver_status,
                "logs": lambda: self.check_system_logs(log_age, log_size), "hardware": self.check_hardware_info,
            }
            if checks:
                for check_name in checks:
                    if check_name in available_checks: self.results.append(available_checks[check_name]())
            else: self.run_all_checks(log_age, log_size, auto_install_pkg)
            report_file = self.generate_report()
            package_file = self.create_package()
            self.logger.info(f"DCU性能分析完成! 报告: {report_file}, 包: {package_file}")
            return package_file
        except Exception as e:
            self.logger.error(f"分析过程失败: {str(e)}", exc_info=self.debug)
            raise

def create_cli():
    parser = argparse.ArgumentParser(description="DCU Performance Analyzer", formatter_class=argparse.RawDescriptionHelpFormatter, epilog="示例:\n  %(prog)s -o /tmp/analysis\n  %(prog)s -c system pcie\n  %(prog)s --auto-install-pkg  # 自动安装缺失的软件包\n  %(prog)s -c pkg --auto-install-pkg  # 只检查并安装软件包")
    parser.add_argument("-c", "--checks", nargs="+", choices=["system", "pcie", "driver", "logs", "hardware", "pkg"], help="指定要运行的检查模块")
    parser.add_argument("-o", "--output", help="指定输出目录")
    parser.add_argument("-t", "--log-age", type=int, default=24, help="收集日志的时间范围(小时) (默认: 24)")
    parser.add_argument("-s", "--log-size", type=int, default=10, help="日志文件大小限制(MB) (默认: 10)")
    parser.add_argument("--auto-install-pkg", action="store_true", help="自动安装缺失的软件包")
    parser.add_argument("-d", "--debug", action="store_true", help="启用调试模式")
    parser.add_argument("-q", "--quiet", action="store_true", help="静默模式")
    parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")
    return parser

def main():
    parser = create_cli()
    args = parser.parse_args()
    try:
        analyzer = DCUPerformanceAnalyzer(output_dir=args.output, debug=args.debug, quiet=args.quiet)
        package_file = analyzer.run_analysis(checks=args.checks, log_age=args.log_age, log_size=args.log_size, auto_install_pkg=args.auto_install_pkg)
        print(f"\n分析完成! 结果包: {package_file}")
        if analyzer.results:
            status_counts = {"PASS": 0, "FAIL": 0, "WARNING": 0, "INFO": 0}
            for result in analyzer.results: status_counts[result.status] += 1
            print("\n检查结果统计:")
            for status, count in status_counts.items(): print(f"  {status}: {count}")
        return 0
    except Exception as e:
        print(f"错误: {str(e)}")
        if args.debug: import traceback; traceback.print_exc()
        return 1

if __name__ == "__main__":
    sys.exit(main())