"""Download CIFAR-10 and extract images for training + FID evaluation.

Usage:
    python prepare_cifar10.py [--data_root ./cifar10_data]

Creates:
    <data_root>/cifar-10-batches-py/   (pickle files used by data.py)
    <data_root>/img/                   (PNG images for FID computation)
"""
import argparse
import os
import pickle
import shutil
import tarfile
import urllib.request

import numpy as np
from PIL import Image

CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"


def is_valid_tar_gz(path: str) -> bool:
    if not os.path.isfile(path):
        return False
    try:
        with tarfile.open(path, "r:gz") as tf:
            return tf.getmember("cifar-10-batches-py") is not None
    except Exception:
        return False


def ensure_clean_tarball(path: str):
    if os.path.isfile(path) and not is_valid_tar_gz(path):
        print(f"Corrupted archive detected, removing: {path}")
        os.remove(path)


def has_complete_batches(batch_dir: str) -> bool:
    required = [
        "data_batch_1",
        "data_batch_2",
        "data_batch_3",
        "data_batch_4",
        "data_batch_5",
        "test_batch",
        "batches.meta",
    ]
    return all(os.path.isfile(os.path.join(batch_dir, name)) for name in required)


def download_cifar10(data_root: str):
    tar_path = os.path.join(data_root, "cifar-10-python.tar.gz")
    batch_dir = os.path.join(data_root, "cifar-10-batches-py")
    if os.path.isdir(batch_dir) and has_complete_batches(batch_dir):
        print(f"Already exists: {batch_dir}")
        return
    if os.path.isdir(batch_dir):
        print(f"Incomplete batch directory detected, removing: {batch_dir}")
        shutil.rmtree(batch_dir, ignore_errors=True)
    os.makedirs(data_root, exist_ok=True)
    ensure_clean_tarball(tar_path)
    if not os.path.isfile(tar_path):
        print(f"Downloading CIFAR-10 to {tar_path} ...")
        urllib.request.urlretrieve(CIFAR10_URL, tar_path)
        print("Download complete.")
        ensure_clean_tarball(tar_path)
    print("Extracting ...")
    try:
        with tarfile.open(tar_path, "r:gz") as tf:
            tf.extractall(data_root)
    except Exception as e:
        print(f"Extraction failed ({type(e).__name__}); re-downloading archive once ...")
        if os.path.isdir(batch_dir):
            shutil.rmtree(batch_dir, ignore_errors=True)
        if os.path.isfile(tar_path):
            os.remove(tar_path)
        urllib.request.urlretrieve(CIFAR10_URL, tar_path)
        with tarfile.open(tar_path, "r:gz") as tf:
            tf.extractall(data_root)
    print(f"Extracted to {batch_dir}")


def extract_images(data_root: str):
    """Extract all training images as PNGs into <data_root>/img/ for FID."""
    img_dir = os.path.join(data_root, "img")
    if os.path.isdir(img_dir) and len(os.listdir(img_dir)) >= 45000:
        print(f"Image folder already populated: {img_dir}")
        return
    os.makedirs(img_dir, exist_ok=True)
    batch_dir = os.path.join(data_root, "cifar-10-batches-py")
    idx = 0
    for batch_id in range(1, 6):
        path = os.path.join(batch_dir, f"data_batch_{batch_id}")
        with open(path, "rb") as f:
            batch = pickle.load(f, encoding="bytes")
        images = batch[b"data"].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
        for img_np in images:
            Image.fromarray(img_np).save(
                os.path.join(img_dir, f"{idx:06d}.png"))
            idx += 1
    print(f"Extracted {idx} training images to {img_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_root", type=str, default="./cifar10_data")
    args = parser.parse_args()

    download_cifar10(args.data_root)
    extract_images(args.data_root)
    print("Done. Use --data_root", args.data_root, "when launching training.")