import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from volcenginesdkcore.rest import ApiException
import volcenginesdkcore
import volcenginesdkecs
import volcenginesdkstorageebs

class VolcCloudECS:
    def __init__(self):
        self.ak = os.getenv("VOLC_ACCESS_KEY")
        self.sk = os.getenv("VOLC_SECRET_KEY")
        self.region = os.getenv("VOLC_REGION", "cn-beijing")
        if not self.ak or not self.sk:
            raise EnvironmentError("请设置 VOLC_ACCESS_KEY 和 VOLC_SECRET_KEY 环境变量。")
        self.init_volc_config()
        self.ecs_api = volcenginesdkecs.ECSApi()
        self.ebs_api = volcenginesdkstorageebs.STORAGEEBSApi()

    def init_volc_config(self):
        config = volcenginesdkcore.Configuration()
        config.ak = self.ak
        config.sk = self.sk
        config.region = self.region
        volcenginesdkcore.Configuration.set_default(config)

    def get_total_cpu_and_mem(self):
        """优化CPU和内存获取，增加错误处理和超时"""
        req = volcenginesdkecs.DescribeInstancesRequest()
        req.max_results = 100  # 增加批量大小
        total_cpu = 0
        total_mem = 0
        next_token = None
        request_count = 0
        max_requests = 50  # 限制最大请求次数，避免无限循环

        while request_count < max_requests:
            if next_token:
                req.next_token = next_token
            try:
                resp = self.ecs_api.describe_instances(req)
                if not resp or not resp.instances:
                    break

                for inst in resp.instances:
                    cpu = inst.cpu_options
                    cores = cpu.core_count if cpu else 0
                    threads = cpu.threads_per_core if cpu else 0
                    total_cpu += cores * threads
                    mem_mib = inst.memory_size or 0
                    total_mem += mem_mib / 1024

                next_token = resp.next_token
                request_count += 1

                if not next_token:
                    break
            except ApiException as e:
                print(f"[火山云ECS] 请求出错: {e}")
                break
            except Exception as e:
                print(f"[火山云ECS] 未知错误: {e}")
                break

        return total_cpu, round(total_mem, 2)

    def get_total_disk(self):
        """优化磁盘获取，增加错误处理和分页限制"""
        req = volcenginesdkstorageebs.DescribeVolumesRequest()
        req.page_size = 100  # 保持较大的页面大小
        req.page_number = 1
        total_disk = 0
        max_pages = 100  # 限制最大页数，避免无限循环

        while req.page_number <= max_pages:
            try:
                resp = self.ebs_api.describe_volumes(req)
                if not resp or not resp.volumes:
                    break

                for vol in resp.volumes:
                    total_disk += int(vol.size or 0)

                # 检查是否还有更多页面
                if len(resp.volumes) < req.page_size:
                    break

                req.page_number += 1

            except ApiException as e:
                print(f"[火山云EBS] 请求出错: {e}")
                break
            except Exception as e:
                print(f"[火山云EBS] 未知错误: {e}")
                break

        return total_disk

    def get_total_resources(self):
        try:
            cpu, mem = self.get_total_cpu_and_mem()
            disk = self.get_total_disk()
            return {
                "cpu": cpu,
                "memory": mem,
                "disk": disk,
                "cloud_provider": "火山云"
            }
        except Exception as e:
            print(f"火山云资源统计出错: {e}")
            return {
                "cpu": 0,
                "memory": 0,
                "disk": 0,
                "cloud_provider": "火山云"
            }
