import os, time
import requests
from zipfile import ZipFile

output_dir = r"X:\AI Data\AI Training\Training Data\diffusiondb"
os.makedirs(output_dir, exist_ok=True)

total_parts = 2000
start_part = 1691  
end_part = total_parts  

base_url = "https://huggingface.co/datasets/poloclub/diffusiondb/resolve/main/images"

def is_zip_valid(file_path):
    """Check if a zip file can be opened successfully."""
    try:
        with ZipFile(file_path, 'r') as zf:
            bad_file = zf.testzip()
            return bad_file is None
    except Exception:
        return False

for i in range(start_part, end_part + 1):
    part_name = f"part-{i:06d}.zip"
    out_path = os.path.join(output_dir, part_name)

    if os.path.exists(out_path) and is_zip_valid(out_path):
        print(f"Skipping {part_name}, already downloaded and valid.")
        continue

    url = f"{base_url}/{part_name}"
    print(f"Downloading {part_name} ...")
    try:
        r = requests.get(url, stream=True)
        r.raise_for_status()
        with open(out_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=1024*1024):
                f.write(chunk)
        print(f"Downloaded {part_name}")

        if not is_zip_valid(out_path):
            print(f"Warning: {part_name} appears corrupted, will re-download next time.")
    except Exception as e:
        print(f"Failed to download {part_name}: {e}")


    time.sleep(10)