import os, time
import zipfile
import shutil
from io import BytesIO

DOWNLOAD_DIR = r"X:\AI Data\AI Training\Training Data\diffusiondb"
OUTPUT_BASE  = r"C:\Users\Abzu\Desktop\diffusiondb\extracted data"
START_PART   = 1836
END_PART     = 2000


def extract_zip_from_ram(zip_path, output_folder):
    # Remove existing folder to avoid mixing files
    if os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(output_folder, exist_ok=True)

    # ✅ Load zip into RAM
    print(f"📥 Loading {zip_path} into RAM...")
    with open(zip_path, "rb") as f:
        zip_data = BytesIO(f.read())  # entire zip stored in memory

    # ✅ Open the in-RAM zip
    with zipfile.ZipFile(zip_data, 'r') as zip_ref:
        members = zip_ref.infolist()
        total = len(members)

        for i, member in enumerate(members, 1):
            zip_ref.extract(member, output_folder)
            print(f"[{i}/{total}] Extracting {member.filename}", end="\r")

    print(f"\n✅ Done extracting to {output_folder}")


def main():
    processed = 0
    for part in range(START_PART, END_PART + 1):
        filename = f"part-{part:06d}.zip"
        zip_path = os.path.join(DOWNLOAD_DIR, filename)

        if not os.path.exists(zip_path):
            print(f"⚠️  {zip_path} not found, skipping...")
            continue

        print(f"\n>>> Processing {filename}")
        extract_zip_from_ram(zip_path, os.path.join(OUTPUT_BASE, f"part-{part:06d}"))

        processed += 1
        if processed % 20 == 0:
            print("⏳ Cooling 10s to protect drives...")
            time.sleep(10)


if __name__ == "__main__":
    main()