"""
Batch product image restyler using OpenAI gpt-image-1 (DALL·E 3)
- Parallel execution (configurable)
- Automatic retry on errors
- Auto crop + resize to 1000×1130
- Respects API rate limits
"""

import os, csv, base64, io, time, threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import backoff
from openai import OpenAI
from PIL import Image

# ----------------------------
# CONFIGURATION
# ----------------------------
API_KEY   = "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"
IN_DIR    = "images_in"
OUT_DIR   = "images_out"
CSV_FILE  = "catalog.csv"     # columns: sku,filename
MAX_WORKERS = 1               # safe range: 3–5
SIZE = "1024x1024"            # valid DALL·E dimension
RATE_LIMIT = 12 

client = OpenAI(api_key=API_KEY)
os.makedirs(OUT_DIR, exist_ok=True)
# ----------------------------
# RATE LIMITER (thread-safe)
# ----------------------------
last_call = 0
lock = threading.Lock()
def rate_limited_call():
    """Ensure we send ≤5 requests per minute total"""
    global last_call
    with lock:
        elapsed = time.time() - last_call
        wait = max(0, RATE_LIMIT - elapsed)
        if wait > 0:
            time.sleep(wait)
        last_call = time.time()

# ----------------------------
# CROPPING & RESIZING
# ----------------------------
def crop_resize_to_target(img_bytes, path_out):
    """Crop and resize image to 1000×1130 while preserving aspect"""
    img = Image.open(io.BytesIO(img_bytes))
    w, h = img.size
    target_w, target_h = 1000, 1130

    # Center-crop to roughly match aspect ratio
    aspect_target = target_w / target_h
    aspect_img = w / h
    if aspect_img > aspect_target:
        new_w = int(h * aspect_target)
        left = (w - new_w) // 2
        box = (left, 0, left + new_w, h)
    else:
        new_h = int(w / aspect_target)
        top = (h - new_h) // 2
        box = (0, top, w, top + new_h)
    cropped = img.crop(box)
    resized = cropped.resize((target_w, target_h), Image.LANCZOS)
    resized.save(path_out, quality=95)

# ----------------------------
# RESTYLE FUNCTION WITH RETRIES
# ----------------------------
@backoff.on_exception(backoff.expo, Exception, max_tries=3, jitter=None)
def restyle_image(path_in, path_out):
    """Send one image to OpenAI and save the restyled output"""
    rate_limited_call()
    with open(path_in, "rb") as img:
        result = client.images.edit(
            model="gpt-image-1",
            image=img,
            prompt=(
                "Restyle the uploaded product image with a pure white background and a natural soft shadow underneath."
        "Keep the product identical (no color or shape changes), including all logos, Text on product, brand names, hardware, and metal details — do NOT change, remove, or alter them. And do not change the text on the stripe or any part of the product. "        
        "Maintain the same position, angle, and proportions of the original image."
        "Produce a clean, realistic studio e-commerce packshot."
            ),
            size=SIZE,
            n=1,
            quality="high"
        )
    b64 = result.data[0].b64_json
    img_bytes = base64.b64decode(b64)
    crop_resize_to_target(img_bytes, path_out)

# ----------------------------
# PER-IMAGE WRAPPER
# ----------------------------
def process_row(row):
    sku = row["sku"].strip()
    fn  = row["filename"].strip()
    inp = os.path.join(IN_DIR, fn)
    out = os.path.join(OUT_DIR, f"{sku}-d.jpg")

    if not os.path.exists(inp):
        return f"{sku}: MISSING INPUT"
    if os.path.exists(out):
        return f"{sku}: SKIP"

    try:
        restyle_image(inp, out)
        return f"{sku}: OK"
    except Exception as e:
        return f"{sku}: ERR {e}"

# ----------------------------
# MAIN EXECUTION
# ----------------------------
if __name__ == "__main__":
    start = time.time()
    with open(CSV_FILE) as f:
        rows = list(csv.DictReader(f))

    total = len(rows)
    print(f"Processing {total} images using {MAX_WORKERS} threads...\n")

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(process_row, row): row for row in rows}
        for i, future in enumerate(as_completed(futures), 1):
            print(f"[{i}/{total}] {future.result()}")

    mins = (time.time() - start) / 60
    print(f"\n✅ Completed {total} images in {mins:.1f} minutes.")