from __future__ import annotations

import argparse
import io
import json
import os
import random
import zipfile
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import requests
from PIL import Image

ANN_URL = "http://images.cocodataset.org/annotations/annotations_trainval2014.zip"

def csr(img, size):
    img = img.convert("RGB")
    w, h = img.size
    s = min(w, h)
    l, t = (w - s) // 2, (h - s) // 2
    return np.asarray(img.crop((l, t, l + s, t + s)).resize((size, size), Image.BICUBIC), dtype=np.uint8)

def fetch_one(item, size):
    url, cap = item
    for _ in range(3):
        try:
            r = requests.get(url, timeout=15)
            if r.status_code == 200:
                return csr(Image.open(io.BytesIO(r.content)), size), cap
        except Exception:
            pass
    return None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--work", default="/root/v6cache")
    ap.add_argument("--n", type=int, default=5000)
    ap.add_argument("--size", type=int, default=256)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--workers", type=int, default=48)
    args = ap.parse_args()
    os.makedirs(args.work, exist_ok=True)

    out_path = os.path.join(args.work, "eval_256.npz")
    if os.path.exists(out_path):
        print(f"[eval-set] already exists at {out_path}", flush=True)
        return

    ann_path = os.path.join(args.work, "captions_val2014.json")
    if not os.path.exists(ann_path):
        print("[eval-set] downloading annotations", flush=True)
        z = os.path.join(args.work, "ann.zip")
        with requests.get(ANN_URL, stream=True, timeout=120) as r:
            with open(z, "wb") as f:
                for chunk in r.iter_content(1 << 20):
                    f.write(chunk)
        with zipfile.ZipFile(z) as zf:
            with zf.open("annotations/captions_val2014.json") as src, open(ann_path, "wb") as dst:
                dst.write(src.read())
        os.remove(z)

    ann = json.load(open(ann_path))
    url_by_id = {im["id"]: im["coco_url"] for im in ann["images"]}
    cap_by_id = {}
    for a in ann["annotations"]:
        cap_by_id.setdefault(a["image_id"], a["caption"])
    items = [(url_by_id[i], cap_by_id[i]) for i in cap_by_id if i in url_by_id]
    random.Random(args.seed).shuffle(items)
    print(f"[eval-set] {len(items)} val2014 pairs available, target {args.n}", flush=True)

    imgs, caps = [], []
    pool = ThreadPoolExecutor(max_workers=args.workers)
    idx, batch = 0, 64
    while len(imgs) < args.n and idx < len(items):
        chunk = items[idx:idx + batch]
        idx += batch
        results = [r for r in pool.map(lambda it: fetch_one(it, args.size), chunk) if r is not None]
        for a, c in results:
            imgs.append(a); caps.append(c)
        if idx % (batch * 20) == 0:
            print(f"[eval-set] {len(imgs)}/{args.n}", flush=True)

    imgs = np.stack(imgs[:args.n])
    caps = np.array(caps[:args.n], dtype=object)
    np.savez(out_path, images=imgs, captions=caps)
    print(f"[eval-set] DONE {imgs.shape} -> {out_path}", flush=True)

if __name__ == "__main__":
    main()