import requests
import re
import ast
import json  # 🌟 加上这一行

def calculate_dynamic_size():
    # ==========================================
    # 1. 真实链接配置 (使用 raw 原始文件下载链接)
    # ==========================================
    py_url = "https://huggingface.co/spaces/waveshaker/universal-lora-studio/resolve/main/model_config.py"
    json_url = "https://huggingface.co/waveshaker/comfy_models/resolve/main/studio_config.json"

    print("🔍 1. 正在从 Hugging Face 下载并解析 model_config.py ...")
    try:
        py_res = requests.get(py_url, timeout=15)
        py_res.raise_for_status()
        py_text = py_res.text
        tree = ast.parse(py_text)
        local_config = {}
        for node in ast.walk(tree):
            if isinstance(node, ast.Assign) and getattr(node.targets[0], 'id', '') == 'MODEL_CONFIG':
                local_config = ast.literal_eval(node.value)
                break
    except Exception as e:
        print(f"❌ 解析 model_config.py 失败: {e}")
        return

    print("🔍 2. 正在从 Hugging Face 下载并解析 studio_config.json ...")
    try:
        json_res = requests.get(json_url, timeout=15)
        json_res.raise_for_status()
        json_str = json_res.text
        # 🌟 容错修复：你截图里的 JSON 在 krea2 的 unets 里少了一个逗号，这里自动补上
        json_str = re.sub(r'\}\s*\{', '}, {', json_str)
        remote_config = json.loads(json_str)
    except Exception as e:
        print(f"⚠️ 解析 studio_config.json 失败: {e}，将仅使用本地兜底配置。")
        remote_config = {}

    # ==========================================
    # 2. 提取并去重 (利用 Set)
    # ==========================================
    unique_files = set()
    def add_file(repo, subfolder, filename):
        if repo and filename:
            clean_sub = subfolder.strip("/") if subfolder else ""
            unique_files.add((repo, clean_sub, filename))

    # 提取 py 配置
    for model_name, model_data in local_config.get("models", {}).items():
        # 1. TE & VAE
        for repo, path in model_data.get("companion_downloads", []):
            parts = path.split("/")
            add_file(repo, "/".join(parts[:-1]), parts[-1])
        # 2. 兜底 UNet
        for name, info in model_data.get("fallback_unets", {}).items():
            add_file(info.get("repo"), info.get("subfolder"), info.get("filename"))
        # 3. 兜底 LoRA
        for name, info in model_data.get("fallback_base_loras", {}).items():
            add_file(info.get("repo"), info.get("subfolder"), info.get("filename"))

    # 提取远程 JSON
    for root_key, data in remote_config.items():
        for category in ["unets", "base_loras"]:
            for item in data.get(category, []):
                add_file(item.get("repo"), item.get("subfolder"), item.get("filename"))

    print(f"\n🧹 去重后共发现 {len(unique_files)} 个独立文件。开始极速计算尺寸...\n")

    # ==========================================
    # 3. 极速查询 API 并计算
    # ==========================================
    total_bytes = 0
    folder_cache = {}

    def get_file_size(repo, subfolder, filename):
        cache_key = f"{repo}/{subfolder}"
        if cache_key not in folder_cache:
            api_url = f"https://huggingface.co/api/models/{repo}/tree/main/{subfolder}"
            try:
                res = requests.get(api_url, timeout=10)
                if res.status_code == 200:
                    folder_cache[cache_key] = {
                        item['path'].split('/')[-1]: item.get('size', 0) 
                        for item in res.json()
                    }
                else:
                    folder_cache[cache_key] = {}
            except:
                folder_cache[cache_key] = {}
        return folder_cache[cache_key].get(filename, 0)

    for repo, subfolder, filename in sorted(unique_files):
        size = get_file_size(repo, subfolder, filename)
        if size > 0:
            total_bytes += size
            print(f"✅ {repo:<25} | {filename[:45]:<45} | {size / (1024**3):>5.2f} GB")
        else:
            print(f"❌ {repo:<25} | {filename[:45]:<45} | 获取失败")

    total_gb = total_bytes / (1024 ** 3)
    print(f"\n" + "="*70)
    print(f"📊 动态计算完成！")
    print(f"💾 去重后文件数: {len(unique_files)}")
    print(f"💾 总计需要下载: {total_gb:.2f} GB")
    print(f"="*70)

if __name__ == "__main__":
    calculate_dynamic_size()