#!/usr/bin/env python3
"""
重试下载失败的图片（429限速后等待重试）
"""

import argparse
import sys
import time
from pathlib import Path
from urllib.parse import urlparse, parse_qs

import requests
from tqdm import tqdm


def parse_urls_file(file_path: Path) -> list:
    """解析URL文件，返回URL和类别信息"""
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    items = []
    for line in lines:
        line = line.strip()
        if not line:
            continue

        parsed = urlparse(line)
        query = parse_qs(parsed.query)
        path = query.get('path', [''])[0]
        filename = query.get('fileName', ['unknown.jpg'])[0]

        parts = path.split('/')
        if len(parts) >= 3:
            category = parts[-2]
        else:
            category = "Unknown"

        items.append({
            'url': line,
            'filename': filename,
            'category': category
        })

    return items


def create_directories(base_dir: Path, categories: set) -> dict:
    dirs = {}
    for cat in categories:
        cat_dir = base_dir / cat
        cat_dir.mkdir(parents=True, exist_ok=True)
        dirs[cat] = cat_dir
    return dirs


def download_file(url: str, output_path: Path, session: requests.Session) -> bool:
    try:
        response = session.get(url, timeout=30, stream=True)
        response.raise_for_status()

        with open(output_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return True
    except Exception as e:
        return False


def get_missing_files(urls_file: Path, output_dir: Path) -> list:
    """获取缺失文件的URL列表"""
    items = parse_urls_file(urls_file)
    missing = []

    for item in items:
        category_dir = output_dir / item['category']
        output_path = category_dir / item['filename']

        if not output_path.exists() or output_path.stat().st_size == 0:
            missing.append(item)

    return missing


def retry_download(urls_file: Path, output_dir: Path, max_retries: int = 3,
                   wait_seconds: int = 60) -> None:
    """重试下载失败的文件"""
    print("=" * 50)
    print("重试下载失败的文件")
    print("=" * 50)

    # 找到缺失的文件
    missing = get_missing_files(urls_file, output_dir)
    print(f"缺失 {len(missing)} 个文件")

    if not missing:
        print("没有缺失文件，下载完成！")
        return

    categories = set(item['category'] for item in missing)
    dirs = create_directories(output_dir, categories)

    session = requests.Session()
    session.headers.update({
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
    })

    success_count = 0
    fail_count = 0

    print(f"\n等待 {wait_seconds} 秒以避免限速...")
    time.sleep(wait_seconds)

    print("\n开始重试下载...")
    for item in tqdm(missing, desc="重试进度"):
        category_dir = dirs[item['category']]
        output_path = category_dir / item['filename']

        if download_file(item['url'], output_path, session):
            success_count += 1
        else:
            fail_count += 1
            time.sleep(1)  # 失败后稍作等待

        time.sleep(0.2)  # 避免过快请求

    print("\n" + "=" * 50)
    print(f"重试完成!")
    print(f"成功: {success_count}")
    print(f"失败: {fail_count}")
    print("=" * 50)


def main():
    parser = argparse.ArgumentParser(description="重试下载失败的文件")
    parser.add_argument("--urls-file", type=str, default="data/dataurls.txt")
    parser.add_argument("--output-dir", type=str, default="data/raw/images")
    parser.add_argument("--wait", type=int, default=60, help="初始等待秒数")

    args = parser.parse_args()

    urls_file = Path(args.urls_file)
    output_dir = Path(args.output_dir)

    if not urls_file.exists():
        print(f"错误: URL文件不存在: {urls_file}")
        sys.exit(1)

    retry_download(urls_file, output_dir, wait_seconds=args.wait)


if __name__ == "__main__":
    main()