#!/usr/bin/env python3

import gi
import os
import re
import logging
import apt_pkg
gi.require_version('OSTree', '1.0')
from gi.repository import OSTree

# 全局常量定义

OVL_DIR_PRE = "/ostree/pkgs/ovl-"
ETC_LAYER_LOWER = "/etc-ovl/etc-lower/"
ETC_LAYER_UPPER = "/etc-ovl/etc-upper/"

USR_LAYER_LOWER = "/usr-ovl/usr-lower/"
USR_LAYER_UPPER = "/usr-ovl/usr-upper/"

VAR_LAYER_LOWER = "/var-ovl/var-lower/lib/"
VAR_LAYER_UPPER = "/var-ovl/var-upper/lib/"

ovl_var_upper_current = None


class DataCollector:
    """数据收集器类"""

    def __init__(self):
        """初始化sysroot和部署信息"""
        self.sysroot = None
        self.booted_deployment = None
        self.deployments = None
        self._initialize()

    def _initialize(self):
        """初始化sysroot"""
        try:
            # 初始化sysroot
            logging.info("[DataCollector] Starting sysroot initialization...")
            self.sysroot = OSTree.Sysroot.new(None)
            logging.info("[DataCollector] OSTree.Sysroot created successfully")
            self.sysroot.set_mount_namespace_in_use()
            self.sysroot.initialize()
            logging.info("[DataCollector] Sysroot initialized successfully")
            self.sysroot.load()
            logging.debug("[DataCollector] Sysroot loaded successfully")

            # 获取当前启动的部署和所有部署
            self.booted_deployment = self.sysroot.get_booted_deployment()
            self.deployments = self.sysroot.get_deployments()

        except Exception as e:
            logging.error(f"[DataCollector] Initialization failed: {str(e)}")
            self.sysroot = None
            self.booted_deployment = None
            self.deployments = None

    def is_initialized(self):
        """检查是否初始化成功"""
        return self.sysroot is not None and self.booted_deployment is not None

    def get_current_deploy_hash(self):
        """获取当前部署的哈希值（包含部署序列号）"""
        if not self.is_initialized():
            logging.error("[DataCollector] System not properly initialized")
            return None

        try:
            if self.booted_deployment:
                # 获取哈希值和部署序列号
                deployment_hash = self.booted_deployment.get_csum()
                deploy_serial = self.booted_deployment.get_deployserial()

                # 组合成完整的哈希值：哈希值.序列号
                full_hash = f"{deployment_hash}.{deploy_serial}"
                return full_hash
            else:
                logging.error("[DataCollector] No booted deployment found")
                return None

        except Exception as e:
            logging.error(f"[DataCollector] Error: {str(e)}")
            return None

    def get_pending_deploy_path(self):
        """获取即将部署的系统的哈希值和部署路径"""
        if not self.is_initialized():
            logging.error("[DataCollector] System not properly initialized")
            return None

        try:
            if not self.deployments or len(self.deployments) == 0:
                logging.error("[DataCollector] No deployments found")
                return None

            # 获取第一个部署（通常是pending部署）
            pending_deployment = self.deployments[0]

            # 检查是否是pending部署（与当前启动的部署不同）
            if self.booted_deployment and pending_deployment and pending_deployment != self.booted_deployment:
                # 获取哈希值和部署序列号
                deployment_hash = pending_deployment.get_csum()
                deploy_serial = pending_deployment.get_deployserial()

                # 组合成完整的哈希值：哈希值.序列号
                full_hash = f"{deployment_hash}.{deploy_serial}"

                # 获取操作系统名称并构建部署路径
                osname = pending_deployment.get_osname()
                deploy_path = f"/ostree/deploy/{osname}/deploy/{full_hash}/"

                # 返回部署路径
                return deploy_path
            else:
                logging.error("[DataCollector] No pending deployment found")
                return None

        except Exception as e:
            logging.error(f"[DataCollector] Error: {str(e)}")
            return None

    def parse_dpkg_status_file(self, file_path):
        """解析dpkg状态文件，返回包名和版本的字典"""
        packages = {}

        if not os.path.exists(file_path):
            logging.warning(f"[DataCollector] Warning: File not found: {file_path}")
            return packages

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            logging.info(f"[DataCollector] Successfully read file, size: {len(content)} bytes")

            # 使用正则表达式匹配包信息
            # dpkg状态文件格式：每个包以"Package: "开头，包含"Version: "字段
            package_blocks = re.split(r'\n\n', content)

            for block in package_blocks:
                if not block.strip():
                    continue

                package_name = None
                package_version = None

                # 解析包名和版本
                for line in block.split('\n'):
                    if line.startswith('Package: '):
                        package_name = line.replace('Package: ', '').strip()
                    elif line.startswith('Version: '):
                        package_version = line.replace('Version: ', '').strip()

                if package_name and package_version:
                    packages[package_name] = package_version

        except Exception as e:
            logging.error(f"[DataCollector] Error: Failed to parse dpkg status file {file_path}: {str(e)}")

        return packages

    def compare_dpkg_status_files(self, file1_path, file2_path):
        """比较两个dpkg状态文件的差异，返回差异包的字典，包含新旧版本信息"""
        packages1 = self.parse_dpkg_status_file(file1_path)
        packages2 = self.parse_dpkg_status_file(file2_path)

        diff_packages = {}

        # 比较两个文件的包差异（版本不同或包不存在）
        for pkg_name, pkg_version in packages1.items():
            if pkg_name in packages2:
                if packages2[pkg_name] != pkg_version:
                    # 版本不同，显示新版本（file2的版本）
                    diff_packages[pkg_name] = {
                        'old_version': pkg_version,
                        'new_version': packages2[pkg_name]
                    }
            else:
                # file2中不存在该包，说明包被删除了
                diff_packages[pkg_name] = {
                    'old_version': pkg_version,
                    'new_version': None,
                    'status': 'removed'
                }

        # 检查file2中有但file1中没有的包（新增的包）
        for pkg_name, pkg_version in packages2.items():
            if pkg_name not in packages1:
                diff_packages[pkg_name] = {
                    'old_version': None,
                    'new_version': pkg_version,
                    'status': 'added'
                }

        return diff_packages

    def generate_diff_dict(self, hash_value):
        """生成upper层与deploy目录对比的字典，返回包名和两个版本信息"""
        if not hash_value:
            logging.error("[DataCollector] Error: Failed to get hash value, cannot generate diff dictionary")
            return None

        # 获取当前部署的OS名称
        osname = get_data_migration_osname()

        # 构建文件路径
        deploy_dpkg_status = f"/ostree/deploy/{osname}/deploy/{hash_value}/var/lib/dpkg/status"

        # upper层dpkg状态文件路径
        ovl_upper_path = f"{OVL_DIR_PRE}{hash_value}{VAR_LAYER_UPPER}dpkg/status"

        logging.info(f"[DataCollector] Deploy directory dpkg status file: {deploy_dpkg_status}")
        logging.info(f"[DataCollector] Upper layer dpkg status file: {ovl_upper_path}")

        # 检查文件是否存在
        if not os.path.exists(deploy_dpkg_status):
            logging.error(f"[DataCollector] Error: Deploy directory dpkg status file not found: {deploy_dpkg_status}")
            return None

        diff_dict = {}

        # 生成overlay-upper层与当前系统部署目录的差异字典
        if os.path.exists(ovl_upper_path):
            # 分别解析两个dpkg状态文件
            deploy_packages = self.parse_dpkg_status_file(deploy_dpkg_status)
            ovl_upper_packages = self.parse_dpkg_status_file(ovl_upper_path)

            # 获取所有包名的并集
            all_packages = set(deploy_packages.keys()) | set(ovl_upper_packages.keys())

            # 构建差异字典
            for pkg_name in all_packages:
                deploy_version = deploy_packages.get(pkg_name)
                ovl_upper_version = ovl_upper_packages.get(pkg_name)

                # 判断包的状态
                if deploy_version is not None and ovl_upper_version is not None:
                    # 两个系统都有这个包，比较版本
                    if deploy_version != ovl_upper_version:
                        # 版本不同
                        diff_dict[pkg_name] = {
                            'deploy_version': deploy_version,
                            'ovl_upper_version': ovl_upper_version
                        }
                    else:
                        # 版本相同，不记录到差异字典
                        continue
                elif deploy_version is not None and ovl_upper_version is None:
                    # 原系统有但修改后系统没有的包（被删除）
                    diff_dict[pkg_name] = {
                        'deploy_version': deploy_version,
                        'ovl_upper_version': 'delete'
                    }
                elif deploy_version is None and ovl_upper_version is not None:
                    # 原系统没有但修改后系统有的包（新增）
                    diff_dict[pkg_name] = {
                        'deploy_version': 'none',
                        'ovl_upper_version': ovl_upper_version
                    }

            logging.info(f"[DataCollector] The dictionary of differences between the upper layer and the current system deployment directory contains {len(diff_dict)} differences package")
            logging.info(f"[DataCollector] Detailed statistics:")
            logging.info(f"  - Packages with different versions: {len([v for v in diff_dict.values() if v['deploy_version'] != 'none' and v['ovl_upper_version'] != 'delete'])}")
            logging.info(f"  - Newly added packages: {len([v for v in diff_dict.values() if v['deploy_version'] == 'none'])}")
            logging.info(f"  - Packages deleted: {len([v for v in diff_dict.values() if v['ovl_upper_version'] == 'delete'])}")
        else:
            logging.warning(f"[DataCollector] Warning: Upper layer dpkg status file not found: {ovl_upper_path}")
            diff_dict = {}
            logging.info("[DataCollector] The dictionary of differences between the upper layer and the current system deployment directory is empty")

        return diff_dict

    def generate_all_diff_dict(self, diff_dict):
        """生成包含pending部署点版本信息的完整差异字典"""
        if not diff_dict:
            logging.warning("[DataCollector] Warning: Diff dictionary is empty, cannot generate complete diff dictionary")
            return None

        # 获取pending部署路径
        pending_deploy_path = self.get_pending_deploy_path()
        if not pending_deploy_path:
            logging.error("[DataCollector] Error: Failed to get pending deploy path")
            return None

        # 构建pending部署点的dpkg状态文件路径
        pending_dpkg_status = f"{pending_deploy_path}var/lib/dpkg/status"
        logging.info(f"[DataCollector] Pending deploy directory dpkg status file: {pending_dpkg_status}")

        # 检查文件是否存在
        if not os.path.exists(pending_dpkg_status):
            logging.error(f"[DataCollector] Error: Pending deploy directory dpkg status file not found: {pending_dpkg_status}")
            return None

        # 解析pending部署点的dpkg状态文件
        pending_packages = self.parse_dpkg_status_file(pending_dpkg_status)

        all_diff_dict = {}

        # 遍历原始差异字典，添加pending部署点的版本信息
        for pkg_name, pkg_info in diff_dict.items():
            # 获取pending部署点的包版本，如果不存在则显示"none"
            pending_version = pending_packages.get(pkg_name, "none")

            # 创建新的包信息字典，包含pending部署点的版本
            all_diff_dict[pkg_name] = {
                'deploy_version': pkg_info['deploy_version'],
                'ovl_upper_version': pkg_info['ovl_upper_version'],
                'pending_version': pending_version
            }

        logging.info(f"[DataCollector] The complete diff dictionary contains {len(all_diff_dict)} differences package")
        logging.info(f"Pending deploy point statistics:")
        logging.info(f"  - Number of packages that exist: {len([v for v in all_diff_dict.values() if v['pending_version'] != 'none'])}")
        logging.info(f"  - Number of packages that do not exist: {len([v for v in all_diff_dict.values() if v['pending_version'] == 'none'])}")

        return all_diff_dict

    def generate_repack_dict(self, all_diff_dict):
        """根据规则生成需要重新打包的包字典"""
        if not all_diff_dict:
            logging.warning("[DataCollector] Warning: Complete diff dictionary is empty, cannot generate repack dictionary")
            return {} , {}

        repack_new_dict = {}
        repack_old_dict = {}

        # 初始化规则统计计数器
        rule_counts = {i: 0 for i in range(1, 10)}

        # 遍历完整差异字典，应用新规则
        for pkg_name, pkg_info in all_diff_dict.items():
            deploy_version = pkg_info['deploy_version']
            ovl_upper_version = pkg_info['ovl_upper_version']
            pending_version = pkg_info['pending_version']

            # 前置条件判断：deploy_version != 'none' and ovl_upper_version != 'delete'
            if deploy_version != 'none' and ovl_upper_version != 'delete':
                # 规则①-⑤：处理正常情况（deploy和upper都存在）
                deploy_vs_upper = self._version_compare(deploy_version, ovl_upper_version)
                
                if deploy_vs_upper < 0:
                    # deploy_version < ovl_upper_version
                    if pending_version != 'none':
                        pending_vs_upper = self._version_compare(pending_version, ovl_upper_version)
                        if pending_vs_upper < 0:
                            # 规则①：pending_version < ovl_upper_version
                            repack_old_dict[pkg_name] = ovl_upper_version
                            rule_counts[1] += 1
                        elif pending_vs_upper > 0:
                            # 规则②：pending_version > ovl_upper_version
                            repack_new_dict[pkg_name] = pending_version
                            rule_counts[2] += 1
                    elif pending_version == 'none':
                        # 规则③：pending_version为"none"
                        repack_old_dict[pkg_name] = ovl_upper_version
                        rule_counts[3] += 1
                
                elif deploy_vs_upper > 0:
                    # deploy_version > ovl_upper_version
                    if pending_version != 'none':
                        pending_vs_upper = self._version_compare(pending_version, ovl_upper_version)
                        if pending_vs_upper > 0:
                            # 规则④：pending_version > ovl_upper_version
                            repack_new_dict[pkg_name] = pending_version
                            rule_counts[4] += 1
                        elif pending_vs_upper < 0 :
                            # 规则⑤：pending_version为"none"或小于ovl_upper_version
                            repack_old_dict[pkg_name] = ovl_upper_version
                            rule_counts[5] += 1
                    elif pending_version == 'none':
                        # 规则⑤：pending_version为"none"
                        repack_old_dict[pkg_name] = ovl_upper_version
                        rule_counts[5] += 1
            
            # 规则⑥：deploy_version不为"none"，但ovl_upper_version为"delete"，且pending_version不为"none"
            elif deploy_version != 'none' and ovl_upper_version == 'delete' and pending_version != 'none':
                repack_new_dict[pkg_name] = pending_version
                rule_counts[6] += 1
            
            # 规则⑦-⑨：deploy_version为"none"
            elif deploy_version == 'none' and ovl_upper_version != 'delete':
                if pending_version != 'none':
                    ovl_vs_pending = self._version_compare(ovl_upper_version, pending_version)
                    if ovl_vs_pending > 0:
                        # 规则⑦：ovl_upper_version > pending_version
                        repack_old_dict[pkg_name] = ovl_upper_version
                        rule_counts[7] += 1
                    elif ovl_vs_pending < 0:
                        # 规则⑧：ovl_upper_version < pending_version
                        repack_new_dict[pkg_name] = pending_version
                        rule_counts[8] += 1
                elif pending_version == 'none':
                    # 规则⑨：pending_version为"none"
                    repack_old_dict[pkg_name] = ovl_upper_version
                    rule_counts[9] += 1

        logging.info(f"[DataCollector] The repack_new_dict contains {len(repack_new_dict)} packages that need to be repacked")
        logging.info(f"[DataCollector] The repack_old_dict contains {len(repack_old_dict)} packages that need to be repacked")
        logging.info(f"New Rules statistics:")
        for rule_num in range(1, 10):
            logging.info(f"  - Rule {rule_num}: {rule_counts[rule_num]} package(s)")

        return repack_new_dict, repack_old_dict

    def _version_compare(self, version1, version2):
        """比较两个版本号，返回-1(version1<version2), 0(相等), 1(version1>version2)"""
        # 处理特殊字符串
        if version1 == 'none' or version1 == 'delete' or version2 == 'none' or version2 == 'delete':
            # 特殊字符串比较，使用简单的字符串比较
            if version1 == version2:
                return 0
            elif version1 < version2:
                return -1
            else:
                return 1

        """
        使用apt_pkg的版本号比较方法
        返回: -1(version1<version2), 0(相等), 1(version1>version2)
        """

        # 使用apt_pkg的标准版本比较函数
        try:
            # apt_pkg.version_compare(version1, version2)
            result = apt_pkg.version_compare(version1, version2)
            if result < 0:
                return -1
            elif result > 0:
                return 1
            else:
                return 0
        except Exception as e:
            logging.info(f"[DataCollector] Warning: apt_pkg version compare failed for '{version1}' vs '{version2}': {str(e)}")
            # 如果apt_pkg比较失败，回退到自定义方法比较
            def split_version(version):
                parts = []
                current = ''
                for char in version:
                    if char.isdigit():
                        if current and not current[-1].isdigit():
                            parts.append(current)
                            current = ''
                        current += char
                    else:
                        if current and current[-1].isdigit():
                            parts.append(current)
                            current = ''
                        current += char
                if current:
                    parts.append(current)
                return parts

            parts1 = split_version(version1)
            parts2 = split_version(version2)

            # 逐个比较部分
            for i in range(max(len(parts1), len(parts2))):
                p1 = parts1[i] if i < len(parts1) else ''
                p2 = parts2[i] if i < len(parts2) else ''

                # 如果都是数字，转换为数字比较
                if p1.isdigit() and p2.isdigit():
                    num1 = int(p1)
                    num2 = int(p2)
                    if num1 < num2:
                        return -1
                    elif num1 > num2:
                        return 1
                else:
                    # 否则使用字符串比较
                    if p1 < p2:
                        return -1
                    elif p1 > p2:
                        return 1

            return 0

def get_data_migration_osname():
    """获取当前部署的OS名称"""
    try:
        collector = DataCollector()
        if collector.is_initialized() and collector.booted_deployment:
            osname = collector.booted_deployment.get_osname()
            logging.info(f"[DataCollector] Current OS name: {osname}")
            return osname
        else:
            logging.warning("[DataCollector] Failed to get OS name, using default 'kylin'")
            return "kylin"
    except Exception as e:
        logging.error(f"[DataCollector] Error getting OS name: {str(e)}")
        return "kylin"

def get_data_migration_dict():
    """主函数"""
    repack_new_dict = None
    repack_old_dict = None
    try:
        collector = DataCollector()
        hash_value = collector.get_current_deploy_hash()
        if hash_value:
            logging.info(f"[DataCollector] Current deploy hash value: {hash_value}")
            ovl_var_upper_current = OVL_DIR_PRE + hash_value + VAR_LAYER_UPPER
            logging.info(f"[DataCollector] ovl_var_upper_current: {ovl_var_upper_current}")

            # 生成差异字典
            logging.info("[DataCollector] Start generating dpkg status file difference dictionary...")
            diff_dict = collector.generate_diff_dict(hash_value)

            # 打印upper层差异字典
            logging.info("[DataCollector] ===== Upper layer difference dictionary with current system deployment directory ======")
            if diff_dict:
                for pkg_name, pkg_info in diff_dict.items():
                    logging.info(f"{pkg_name}: {pkg_info}")
            else:
                logging.info("[DataCollector] The upper layer difference dictionary with current system deployment directory is empty")

            # 生成包含pending部署点的完整差异字典
            logging.info("[DataCollector] Start generating complete difference dictionary with pending deployment points...")
            all_diff_dict = collector.generate_all_diff_dict(diff_dict)

            # 打印完整差异字典
            logging.info("[DataCollector] ====== Complete difference dictionary with pending deployment points ======")
            if all_diff_dict:
                for pkg_name, pkg_info in all_diff_dict.items():
                    logging.info(f" - {pkg_name}: {pkg_info}")
            else:
                logging.info("[DataCollector] The complete difference dictionary with pending deployment points is empty")

            # 生成重新打包字典
            logging.info("[DataCollector] Start generating package dictionary with pending deployment points that need to be repacked(new)...")
            repack_new_dict, repack_old_dict   = collector.generate_repack_dict(all_diff_dict)

            # 打印重新打包字典（新版本）
            logging.info("[DataCollector] ====== Package dictionary with pending deployment points that need to be repacked(new) ======")
            if repack_new_dict:
                for pkg_name, version in repack_new_dict.items():
                    logging.info(f" - {pkg_name}: {version}")
            else:
                logging.info("[DataCollector] There are no packages that need to be repacked(new) with pending deployment points")

            # 打印重新打包字典（旧版本）
            logging.info("[DataCollector] ====== Package dictionary with pending deployment points that need to be repacked(old) ======")
            if repack_old_dict:
                for pkg_name, version in repack_old_dict.items():
                    logging.info(f" - {pkg_name}: {version}")
            else:
                logging.info("[DataCollector] There are no packages that need to be repacked(old) with pending deployment points")
        else:
            logging.info("[DataCollector] Failed to get current deploy hash value")
    except Exception as e:
        logging.error(f"[DataCollector] Error occurred during data migration scanning: {str(e)}")

    return repack_new_dict, repack_old_dict

# if __name__ == "__main__":
#     # -----------------------------
#     # 配置 logging 打印到终端
#     # -----------------------------
#     logging.basicConfig(
#         level=logging.INFO,  # 日志级别，可改为 DEBUG 查看详细信息
#         format="%(asctime)s [%(levelname)s] %(message)s",
#         datefmt="%Y-%m-%d %H:%M:%S",
#     )
#     # 获取 logger 对象
#     logger = logging.getLogger(__name__)

#     main()