from pyVim.connect import SmartConnect, Disconnect
from pyVmomi import vim
from config.settings import MANAGEMENT_NODES
from utils.logger import logger
from dataclasses import dataclass
from typing import List, Optional, Dict


@dataclass
class Connection:
    """包装连接对象，包含节点信息和ServiceInstance"""
    name: str  # 节点名称
    host: str  # 主机地址
    node_type: str  # vcenter 或 esxi
    si: any  # ServiceInstance
    max_concurrent: int  # 该节点的最大并发数


class VCenterManager:
    """管理多个vCenter/ESXi连接"""

    def __init__(self):
        # 构造方法，创建对象时自动执行
        self.connections: List[Connection] = []  # 属性1：成功连接列表
        self.failed_nodes: List[Dict] = []  # 属性2：失败节点列表

    def connect_all(self) -> List[Connection]:
        """
        连接所有配置的管理节点
        :return: 成功连接的列表
        """
        logger.info(f"开始连接 {len(MANAGEMENT_NODES)} 个管理节点...")

        for node in MANAGEMENT_NODES:
            name = node['name']
            host = node['host']
            user = node['user']
            password = node['password']
            node_type = node['type']
            max_concurrent = node.get('max_delete_concurrent', 4)

            logger.info(f"正在连接节点: {name} ({host}) [{node_type}]")

            si = self._connect_single(host, user, password, node_type)

            if si:
                conn = Connection(
                    name=name,
                    host=host,
                    node_type=node_type,
                    si=si,
                    max_concurrent=max_concurrent
                )
                self.connections.append(conn)
                logger.info(f"✅ {name} 连接成功")
            else:
                self.failed_nodes.append({
                    'name': name,
                    'host': host,
                    'type': node_type
                })
                logger.error(f"❌ {name} 连接失败")

        logger.info(f"连接完成: 成功 {len(self.connections)}/{len(MANAGEMENT_NODES)}")
        return self.connections

    def _connect_single(self, host: str, user: str, password: str, node_type: str):
        """连接单个节点"""
        try:
            si = SmartConnect(
                host=host,
                user=user,
                pwd=password,
                disableSslCertValidation=True
            )
            return si
        except Exception as e:
            logger.error(f"连接 {host} 失败: {str(e)}")
            return None

    def get_connection(self, name: str = None) -> Optional[Connection]:
        """
        获取指定名称的连接，不指定则返回第一个
        """
        if not self.connections:
            return None

        if name:
            for conn in self.connections:
                if conn.name == name:
                    return conn
            return None

        return self.connections[0]

    def get_all_connections(self) -> List[Connection]:
        """获取所有成功连接"""
        return self.connections

    def get_vcenter_connections(self) -> List[Connection]:
        """仅获取 vcenter 类型的连接"""
        return [c for c in self.connections if c.node_type == 'vcenter']

    def get_esxi_connections(self) -> List[Connection]:
        """仅获取 esxi 类型的连接"""
        return [c for c in self.connections if c.node_type == 'esxi']

    def disconnect_all(self):
        """关闭所有连接"""
        for conn in self.connections:
            try:
                Disconnect(conn.si)
                logger.info(f"已断开 {conn.name}")
            except Exception as e:
                logger.warning(f"断开 {conn.name} 时出错: {e}")
        self.connections = []

    def __enter__(self):
        """上下文管理器支持"""
        self.connect_all()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """确保退出时断开连接"""
        self.disconnect_all()


# ==================== 使用示例 ====================

def main():
    # 方式1：手动管理连接
    manager = VCenterManager()
    connections = manager.connect_all()

    # 遍历所有连接执行操作
    for conn in connections:
        print(f"\n处理节点: {conn.name} ({conn.host})")
        print(f"最大并发: {conn.max_concurrent}")

        # 获取内容视图
        content = conn.si.RetrieveContent()

        # 获取所有VM
        from pyVmomi import vim
        obj_view = content.viewManager.CreateContainerView(
            content.rootFolder, [vim.VirtualMachine], True
        )
        vms = obj_view.view
        print(f"该节点有 {len(vms)} 台虚拟机")
        obj_view.Destroy()

    # 断开所有连接
    manager.disconnect_all()


def main_with_context():
    # 方式2：使用上下文管理器（推荐，自动断开）
    with VCenterManager() as manager:
        # 获取所有vcenter连接
        vcenters = manager.get_vcenter_connections()

        for conn in vcenters:
            print(f"处理vCenter: {conn.name}")
            # 执行操作...

        # 退出时自动断开所有连接


def main_single_operation():
    # 方式3：只连接特定节点
    manager = VCenterManager()
    manager.connect_all()

    # 获取指定节点
    prod_vc = manager.get_connection('vcenter-prod')
    if prod_vc:
        print(f"连接到生产环境: {prod_vc.host}")
        # 执行操作...

    manager.disconnect_all()


if __name__ == '__main__':
    main()