import json
import urllib.request
import urllib.error
import sys
from huggingface_hub import HfApi

# ================= 配置区 =================
SPACE_ID = "waveshaker/character_loras_gallery"
MODEL_REPO_ID = "waveshaker/character_loras"
# ==========================================

OUTPUT_FILE = "models.json"

def fetch_url_json(url):
    print(f"正在请求: {url}")
    try:
        req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
        with urllib.request.urlopen(req, timeout=60) as response:
            return json.loads(response.read().decode())
    except urllib.error.HTTPError as e:
        if e.code == 401:
            print(f"\n❌ 401 错误: 访问被拒绝。请确保 Space 和 Model 仓库是 Public (公开) 状态。")
        elif e.code == 404:
            print(f"\n❌ 404 错误: 找不到文件，可能是首次生成。")
        else:
            print(f"\n❌ HTTP 错误 {e.code}")
        return None
    except Exception as e:
        print(f"\n❌ 网络请求失败: {e}")
        return None

def get_all_repo_files(repo_id):
    api = HfApi()
    try:
        all_files = list(api.list_repo_files(repo_id=repo_id, repo_type="model"))
    except Exception as e:
        print(f"❌ 获取文件列表失败: {e}")
        return None
    
    return [f for f in all_files if f.endswith('.safetensors')]

def custom_sort_key(file_path):
    parts = file_path.split('/')
    top_dir = parts[0] if len(parts) > 0 else ""
    sub_dir = parts[1] if len(parts) > 1 else ""

    top_order = {"krea2": 1, "anima": 2, "flux1_dev": 3, "pony": 4, "sd15": 5}
    sub_order = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}

    top_weight = top_order.get(top_dir, 999)
    sub_weight = sub_order.get(sub_dir, 999)

    return (top_weight, sub_weight, file_path)

def main():
    # 1. 获取 Model Repo 里的所有文件
    print("\n--- 步骤 1: 拉取 Model Repo 文件列表 ---")
    remote_files = get_all_repo_files(MODEL_REPO_ID)
    
    if remote_files is None:
        print("拉取 Model Repo 失败，脚本终止。")
        sys.exit(1)
        
    print(f"✅ 找到 {len(remote_files)} 个 safetensors 文件。")

    # 2. 从 Space 获取现有的 models.json
    print("\n--- 步骤 2: 拉取 Space 现有的 models.json ---")
    space_subdomain = SPACE_ID.replace("/", "-").replace("_", "-").lower()
    config_url = f"https://{space_subdomain}.static.hf.space/models.json"
    
    existing_data = fetch_url_json(config_url)
    existing_items = {}
    
    if existing_data:
        print("✅ 成功拉取现有的 models.json！")
        existing_items = existing_data.get("items", {})
        
        if not existing_items and "titles" in existing_data:
            print("⚠️ 检测到旧版 titles 结构，正在自动升级为新版 items 结构...")
            old_titles = existing_data.get("titles", {})
            for k, v in old_titles.items():
                existing_items[k] = {"title": v, "trigger": "", "dataset": ""}
    else:
        print("⚠️ 无法获取现有的 models.json，将从头生成。")

    # 3. 比对与合并
    print("\n--- 步骤 3: 比对并合并 ---")
    merged_items = {}
    
    old_keys = set(existing_items.keys())
    newly_added_paths = [] # 👈 新增：用来记录新加入了哪些路径
    
    for file_path in remote_files:
        if file_path in existing_items:
            merged_items[file_path] = existing_items[file_path]
            old_keys.remove(file_path)
        else:
            # 生成空字典并记录路径
            merged_items[file_path] = {
                "title": "",
                "trigger": "",
                "dataset": ""
            }
            newly_added_paths.append(file_path) 

    deleted_count = len(old_keys)
    if deleted_count > 0:
        print(f"⚠️ 发现 {deleted_count} 个在云端已删除的文件，已剔除。")

    # 4. 按自定义规则排序
    print("正在进行自定义排序...")
    sorted_items = {k: merged_items[k] for k in sorted(merged_items.keys(), key=custom_sort_key)}

    # 5. 保存本地文件
    output_data = {"items": sorted_items}
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

    # 6. 终端输出总结与待填清单
    print(f"\n🎉 处理完成！已保存到本地: {OUTPUT_FILE}")
    print(f"  - 远端真实文件总数: {len(remote_files)}")
    print(f"  - 本地保留的旧数据: {len(remote_files) - len(newly_added_paths)}")
    print(f"  - 本次新增待填写: {len(newly_added_paths)}")
    print(f"  - 清理已失效文件: {deleted_count}")

    # 👇 核心修改点：如果存在新增，直接在终端打印出来
    if newly_added_paths:
        print("\n" + "="*60)
        print("⚠️ 以下新增文件需要你手动填写 title / trigger / dataset：")
        print("="*60)
        for path in sorted(newly_added_paths, key=custom_sort_key):
            print(f"👉 {path}")
        print("="*60)
        print("👉 请打开 models.json，搜索上述路径，将对应的空字符串补齐。然后上传覆盖 Space。")
    else:
        print("\n✅ 没有新增文件，无需手动修改。")

if __name__ == "__main__":
    main()