from __future__ import annotations
import argparse, io, json, os, random, zipfile
from concurrent.futures import ThreadPoolExecutor
import numpy as np
import requests
import torch
from PIL import Image
from diffusers import AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer

SCALE = 0.18215
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

@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--work", default="/root/pm4")
    ap.add_argument("--n-train", type=int, default=85000)
    ap.add_argument("--size", type=int, default=256)
    ap.add_argument("--max-tokens", type=int, default=40)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--workers", type=int, default=48)
    ap.add_argument("--vae", default="stabilityai/sd-vae-ft-mse")
    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
    args = ap.parse_args()
    os.makedirs(args.work, exist_ok=True)
    dev = "cuda"

    ann_path = os.path.join(args.work, "captions_train2014.json")
    if not os.path.exists(ann_path):
        print("[data] 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_train2014.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(0).shuffle(items)
    print(f"[data] {len(items)} train2014 image/caption pairs available; target {args.n_train}", flush=True)

    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
    tok = CLIPTokenizer.from_pretrained(args.clip)
    txt = CLIPTextModel.from_pretrained(args.clip).to(dev).half().eval()

    lat_list, seq_list, pool_list = [], [], []
    pool = ThreadPoolExecutor(max_workers=args.workers)
    got, idx, nb = 0, 0, 0
    print("[data] starting download/encode loop", flush=True)
    while got < args.n_train and idx < len(items):
        chunk = items[idx:idx + args.batch]
        idx += args.batch
        nb += 1
        try:
            results = [r for r in pool.map(lambda it: fetch_one(it, args.size), chunk) if r is not None]
            if not results:
                print(f"[data] batch {nb}: 0 ok (skipped)", flush=True)
                continue
            imgs = np.stack([r[0] for r in results]).astype(np.float32) / 127.5 - 1.0
            caps = [r[1] for r in results]
            x = torch.from_numpy(imgs).permute(0, 3, 1, 2).to(dev).half()
            lat_list.append((vae.encode(x).latent_dist.mean * SCALE).cpu().numpy().astype(np.float16))
            t = tok(caps, padding="max_length", max_length=args.max_tokens, truncation=True, return_tensors="pt").to(dev)
            o = txt(**t)
            seq_list.append(o.last_hidden_state.cpu().numpy().astype(np.float16))
            pool_list.append(o.pooler_output.cpu().numpy().astype(np.float16))
            got += len(results)
        except Exception as e:
            print(f"[data] batch {nb} ERROR {type(e).__name__}: {str(e)[:100]}", flush=True)
            continue
        if nb % 10 == 0:
            print(f"[data] cached {got}/{args.n_train} (batch {nb})", flush=True)

    lat = np.concatenate(lat_list)[:args.n_train]
    seq = np.concatenate(seq_list)[:args.n_train]
    pool_ = np.concatenate(pool_list)[:args.n_train]
    np.save(os.path.join(args.work, "latents.npy"), lat)
    np.save(os.path.join(args.work, "text_seq.npy"), seq)
    np.save(os.path.join(args.work, "text_pool.npy"), pool_)
    t = tok([""], padding="max_length", max_length=args.max_tokens, truncation=True, return_tensors="pt").to(dev)
    o = txt(**t)
    np.save(os.path.join(args.work, "null_seq.npy"), o.last_hidden_state.cpu().numpy().astype(np.float16))
    np.save(os.path.join(args.work, "null_pool.npy"), o.pooler_output.cpu().numpy().astype(np.float16))
    print(f"[data] DONE latents{lat.shape} seq{seq.shape} -> {args.work}", flush=True)

if __name__ == "__main__":
    main()