import json
import os

from openai import OpenAI
from tqdm import tqdm


os.environ["OPENAI_API_KEY"] = (
    "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"
)

IMAGE_DIR = "/opt/datasets/bd/gid/png"
MAP_FILE = "batchapi/image_file_ids.json"

client = OpenAI()

# Load existing mapping (resume-safe)
if os.path.exists(MAP_FILE):
    with open(MAP_FILE, "r") as f:
        image_map = json.load(f)
else:
    image_map = {}

for img in tqdm(os.listdir(IMAGE_DIR)):
    if not img.lower().endswith((".png", ".jpg", ".jpeg")):
        continue

    if img in image_map:
        continue  # already uploaded

    img_path = os.path.join(IMAGE_DIR, img)

    file = client.files.create(file=open(img_path, "rb"), purpose="vision")

    image_map[img] = file.id

    # Save after every upload (crash-safe)
    with open(MAP_FILE, "w") as f:
        json.dump(image_map, f, indent=2)

print("All images uploaded and cached.")