"""下载测试图片到 uploads 目录"""

import os
import requests
from pathlib import Path

def download_test_images():
    """从 Picsum 下载测试图片"""
    
    # 确保目录存在
    base_dir = Path("uploads")
    dirs = ["thumbnail", "watermarked", "original"]
    
    for dir_name in dirs:
        (base_dir / dir_name).mkdir(parents=True, exist_ok=True)
    
    # 下载图片（与 seed_data.py 中的作品ID对应）
    # 使用不同的 seed 值获得不同的图片
    images = [
        {"id": 1, "seed": 134, "title": "活动海报"},
        {"id": 2, "seed": 573, "title": "马年招聘海报"},
        {"id": 3, "seed": 60, "title": "新中式主视觉"},
        {"id": 4, "seed": 237, "title": "品牌设计"},
        {"id": 5, "seed": 456, "title": "电商Banner"},
        {"id": 6, "seed": 789, "title": "节日海报"},
        {"id": 7, "seed": 321, "title": "UI设计"},
        {"id": 8, "seed": 654, "title": "插画作品"},
        {"id": 9, "seed": 987, "title": "摄影作品"},
        {"id": 10, "seed": 147, "title": "Logo设计"},
        {"id": 11, "seed": 258, "title": "包装设计"},
        {"id": 12, "seed": 369, "title": "名片设计"},
        {"id": 13, "seed": 741, "title": "宣传单设计"},
        {"id": 14, "seed": 852, "title": "画册设计"},
        {"id": 15, "seed": 963, "title": "展板设计"},
    ]
    
    print("[*] 开始下载测试图片...\n")
    
    for img in images:
        img_id = img["id"]
        seed = img["seed"]
        title = img["title"]
        
        print(f"[+] 下载 [{title}] ...")
        
        # 下载缩略图 (400x300)
        thumbnail_url = f"https://picsum.photos/seed/{seed}/400/300"
        download_image(thumbnail_url, base_dir / "thumbnail" / f"{img_id}.jpg")
        
        # 下载水印图 (800x600)
        watermarked_url = f"https://picsum.photos/seed/{seed}/800/600"
        download_image(watermarked_url, base_dir / "watermarked" / f"{img_id}.jpg")
        
        # 下载原图 (1920x1080)
        original_url = f"https://picsum.photos/seed/{seed}/1920/1080"
        download_image(original_url, base_dir / "original" / f"{img_id}.jpg")
        
        print(f"    [OK] {title} 下载完成\n")
    
    print("[*] 所有测试图片下载完成！")
    print(f"\n图片保存路径: {base_dir.absolute()}")

def download_image(url, save_path):
    """下载单个图片"""
    try:
        # 设置超时时间，避免卡住
        response = requests.get(url, timeout=30, allow_redirects=True)
        response.raise_for_status()
        
        # 保存图片
        with open(save_path, 'wb') as f:
            f.write(response.content)
        
        # 显示文件大小
        size_kb = len(response.content) / 1024
        print(f"    |- {save_path.name}: {size_kb:.1f} KB")
        
    except Exception as e:
        print(f"    [ERROR] 下载失败: {e}")

def clear_images():
    """清空所有测试图片"""
    base_dir = Path("uploads")
    dirs = ["thumbnail", "watermarked", "original"]
    
    print("[*] 清空测试图片...\n")
    
    count = 0
    for dir_name in dirs:
        dir_path = base_dir / dir_name
        if dir_path.exists():
            for file in dir_path.glob("*.jpg"):
                file.unlink()
                count += 1
                print(f"    删除: {file}")
    
    print(f"\n[OK] 已删除 {count} 个图片文件")

if __name__ == "__main__":
    import sys
    
    # 检查命令行参数
    if len(sys.argv) > 1 and sys.argv[1] == "clear":
        clear_images()
    else:
        download_test_images()
