import yaml
import os, time
from pyVmomi import vim
from pyVim.connect import SmartConnect, Disconnect
from config.settings import MANAGEMENT_NODES,YAML_OUTPUT_PATH
from utils.logger import logger


def connect_vcenter(host):
    """根据 NodeHost 连接对应 vCenter"""
    for node in MANAGEMENT_NODES:
        if node['host'] == host:
            return SmartConnect(
                host=node['host'],
                user=node['user'],
                pwd=node['password'],
                disableSslCertValidation=True
            )
    raise Exception(f"未找到 {host} 的连接信息")


def main(dele_snapshots):
    """根据 YAML 信息删除旧快照"""
    # print(dele_snapshots)
    if not dele_snapshots:
        logger.info("没有需要删除的快照")
        return

    # 按 NodeHost 分组，避免重复连接
    grouped = {}
    for snapshot in dele_snapshots:
        if snapshot['is_old']:
            grouped.setdefault(snapshot['NodeHost'], []).append(snapshot)

    for host, snapshots in grouped.items():
        logger.info(f"连接到 {host} 删除快照...")
        si = connect_vcenter(host)

        try:
            for snapshot in snapshots:
                delete_snapshot(si, snapshot)  # 调用删除快照函数
        finally:
            Disconnect(si)


# 执行快照删除的（核心）函数
def delete_snapshot(si, snapshot_info):
    """执行快照删除"""
    content = si.RetrieveContent()
    snap_name = f"{snapshot_info['VMName']}-{snapshot_info['Snapshot Name']}-({snapshot_info['MOID']})"

    vm = find_vm_by_name(content, snapshot_info['VMName'])  # VMName即根据获取到的虚拟机名称查找虚拟机是否存放

    if not vm:
        logger.info(f"未找到 VM: {snapshot_info['VMName']}")
        return
    # 检查该虚拟机是否有快照
    if not vm.snapshot:
        logger.warning(snap_name,"：快照不存在")
        return

    # 虚拟机的快照列表中找到具有指定 MOID 的快照对象
    snapshot_obj = find_snapshot_by_moid(
        vm.snapshot.rootSnapshotList,
        snapshot_info['MOID']
    )
    if not snapshot_obj:
        logger.warning(snap_name,"：未找到")
        return

    logger.info(f"正在删除 Snapshot: {snap_name}")
    try:

        """删除快照核心代码，调用快照对象的 RemoveSnapshot_Task 方法执行。removeChildren = False：表示删除该快照时不删除其子快照。"""
        task = snapshot_obj.RemoveSnapshot_Task(removeChildren=False)

        # 改进的等待逻辑
        while task.info.state in [vim.TaskInfo.State.running, vim.TaskInfo.State.queued]:
            time.sleep(1)  # 避免 CPU 空转，每秒检查一次

        # 更完整的状态判断
        if task.info.state == vim.TaskInfo.State.success:
            logger.info(f"✅ 删除成功: {snap_name}")
        elif task.info.state == vim.TaskInfo.State.error:
            error_msg = task.info.error.msg if task.info.error else "未知错误"
            logger.error(f"❌ 删除失败: {snap_name}, 错误: {error_msg}")

        else:
            # 处理其他状态（如 cancelled）
            logger.warning(f"⚠️ 任务未成功完成，状态: {task.info.state}")

    except Exception as e:
        logger.exception(f"删除快照时发生异常: {snap_name}")
        raise


def find_vm_by_name(content, vm_name):
    """根据 VM 名称查找虚拟机对象"""
    container = content.viewManager.CreateContainerView(
        content.rootFolder, [vim.VirtualMachine], True
    )

    for vm in container.view:
        if vm.name == vm_name:
            container.Destroy()
            return vm

    container.Destroy()
    return None


def find_snapshot_by_moid(snapshot_tree, moid):
    """递归查找 snapshot 对象"""
    for snapshot in snapshot_tree:
        if snapshot.snapshot._moId == moid:
            return snapshot.snapshot

        if snapshot.childSnapshotList:
            result = find_snapshot_by_moid(snapshot.childSnapshotList, moid)
            if result:
                return result
    return None


def load_old_snapshots(file_path):
    """从 YAML 文件中加载旧快照"""
    if os.path.exists(file_path):
        with open(file_path, 'r', encoding='utf-8') as yaml_file:
            return yaml.safe_load(yaml_file)
    else:
        logger.error(f"{file_path} 文件不存在.")
        return []


if __name__ == "__main__":
    old_snapshots = load_old_snapshots(YAML_OUTPUT_PATH)
    main(old_snapshots)