#!/usr/bin/env python3
"""Analisis gambar LEIT (27-77) paralel + rotasi model — caption IG/FB 2 paragraf + CTA + hashtag.
Progress disimpan per gambar di leit_ig_40/captions/ (aman di-restart)."""
import base64, json, io, os, sys, time, threading, urllib.request, urllib.error
from PIL import Image
from concurrent.futures import ThreadPoolExecutor, as_completed

ROOT = "leit_ig_40"
OUT_DIR = os.path.join(ROOT, "captions")
os.makedirs(OUT_DIR, exist_ok=True)

key = open("/home/mantapman/.openrouter_key").read().strip()

MODELS = [
    "google/gemma-4-26b-a4b-it:free",
    "nvidia/nemotron-nano-12b-v2-vl:free",
    "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free",
    "google/gemma-4-31b-it:free",
]

PROMPT = """Analisis gambar ini untuk postingan Instagram/Facebook LEIT (Leadership-English Institute — kursus bahasa Inggris di Surabaya, tagline "English Excellence for a Better Future").

Berikan dalam Bahasa Indonesia dengan format EXACT berikut:

[DESKRIPSI]
Deskripsi singkat isi gambar: apa yang ditampilkan, semua teks yang terbaca di gambar, dan tema/topiknya.

[CAPTION]
Caption 2 paragraf (Bahasa Indonesia, gaya Facebook/Instagram — humanis, hangat, natural). Paragraf 1: menarik perhatian, nyambung dengan isi/tema gambar. Paragraf 2: penutup + CTA (ajakan daftar/kontak). Boleh singgung program LEIT: speaking & semua skill dilatih, kelas kecil pembinaan maksimal, hasil terukur. Kontak: WA Surabaya Timur 0811 331 700, Surabaya Barat 0811 396 488. JANGAN mulai dengan "Terima kasih". Langsung ke inti. Maksimal ~500 karakter total.

[HASHTAG]
6-8 hashtag relevan, campur #LEIT #LEITSurabaya #KursusBahasaInggris dengan tag tematik sesuai gambar."""

# Kumpulkan semua gambar berurutan
files = []
for folder in ["27-40", "41-50", "51-60", "61-77"]:
    d = os.path.join(ROOT, folder)
    if not os.path.isdir(d):
        continue
    for fn in sorted(os.listdir(d), key=lambda f: int(os.path.splitext(f)[0])):
        if fn.endswith(".png"):
            files.append((int(os.path.splitext(fn)[0]), os.path.join(d, fn)))

files.sort(key=lambda x: x[0])
print(f"Total {len(files)} gambar", flush=True)

def analyze(num, path):
    out_file = os.path.join(OUT_DIR, f"{num}.json")
    if os.path.exists(out_file):
        return num, "SKIP (sudah ada)"
    im = Image.open(path).convert("RGB")
    im.thumbnail((640, 640))
    buf = io.BytesIO()
    im.save(buf, "JPEG", quality=78)
    b64 = base64.b64encode(buf.getvalue()).decode()
    last_err = ""
    for mi, model in enumerate(MODELS):
        for attempt in range(3):
            payload = {
                "model": model,
                "messages": [{"role": "user", "content": [
                    {"type": "text", "text": PROMPT},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}
                ]}],
                "max_tokens": 700
            }
            req = urllib.request.Request(
                "https://openrouter.ai/api/v1/chat/completions",
                data=json.dumps(payload).encode(),
                headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"}
            )
            try:
                with urllib.request.urlopen(req, timeout=150) as resp:
                    data = json.loads(resp.read().decode())
                    text = data["choices"][0]["message"]["content"]
                with open(out_file, "w") as f:
                    json.dump({"num": num, "file": os.path.basename(path), "analysis": text}, f, ensure_ascii=False)
                return num, f"OK ({len(text)} chars, model={model})"
            except urllib.error.HTTPError as e:
                body = e.read().decode()[:150]
                last_err = f"HTTP {e.code} {body}"
                if e.code == 429:
                    time.sleep(8 + attempt * 10 + mi * 3)
                else:
                    break
            except Exception as e:
                last_err = str(e)
                time.sleep(5)
    with open(out_file, "w") as f:
        json.dump({"num": num, "file": os.path.basename(path), "analysis": f"ERROR: {last_err}"}, f, ensure_ascii=False)
    return num, f"FAIL: {last_err}"

done = [0]
lock = threading.Lock()

def worker(item):
    num, path = item
    res = analyze(num, path)
    with lock:
        done[0] += 1
        print(f"[{done[0]}/{len(files)}] {res[0]}.png -> {res[1]}", flush=True)
    return res

with ThreadPoolExecutor(max_workers=3) as ex:
    futures = [ex.submit(worker, it) for it in files]
    for fut in as_completed(futures):
        fut.result()

print("\nSELESAI semua analisis!")