#!/usr/bin/env python3
"""
Generate plate images with Replicate (black-forest-labs/flux-1.1-pro by default).

Reusable across issues. Reads a prompts JSON file — an array of objects with at
least `slug` and `prompt` — and writes one `<slug>.png` per entry.

  python3 scripts/generate-images.py --prompts data/prompts/issue-02.json --issue 02

Output layout: <out>/<slug>/<slug>.png  (a Hugo page-bundle dir per plate).
With --issue NN and no --out, <out> defaults to content/posts/NN.

The Replicate token is read from $REPLICATE_API_TOKEN, or from ~/.env if unset.
Only stdlib is used (no `requests`, no `replicate`).
"""
import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

API_ROOT = "https://api.replicate.com/v1"
SITE = Path(__file__).resolve().parent.parent
TERMINAL = {"succeeded", "failed", "canceled"}


def load_token() -> str:
    tok = os.environ.get("REPLICATE_API_TOKEN", "").strip()
    if tok:
        return tok
    env = Path.home() / ".env"
    if env.is_file():
        for line in env.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, _, val = line.partition("=")
            if key.strip() == "REPLICATE_API_TOKEN":
                return val.strip().strip("'").strip('"')
    sys.exit("REPLICATE_API_TOKEN not found in environment or ~/.env")


def _req(url: str, token: str, method: str = "GET", body: dict | None = None,
         extra_headers: dict | None = None):
    data = json.dumps(body).encode() if body is not None else None
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    if extra_headers:
        headers.update(extra_headers)
    r = urllib.request.Request(url, data=data, headers=headers, method=method)
    with urllib.request.urlopen(r, timeout=120) as resp:
        return json.loads(resp.read().decode())


def create_prediction(model: str, inp: dict, token: str, retries: int = 3) -> dict:
    url = f"{API_ROOT}/models/{model}/predictions"
    for attempt in range(retries):
        try:
            return _req(url, token, "POST", {"input": inp}, {"Prefer": "wait"})
        except urllib.error.HTTPError as e:
            if e.code in (429, 500, 502, 503, 504) and attempt < retries - 1:
                time.sleep(2 ** attempt * 3)
                continue
            detail = e.read().decode(errors="replace")
            raise RuntimeError(f"HTTP {e.code} creating prediction: {detail}") from e
        except urllib.error.URLError as e:
            if attempt < retries - 1:
                time.sleep(2 ** attempt * 3)
                continue
            raise RuntimeError(f"network error creating prediction: {e}") from e
    raise RuntimeError("exhausted retries creating prediction")


def wait_for(pred: dict, token: str, timeout_s: int = 300) -> dict:
    deadline = time.time() + timeout_s
    while pred.get("status") not in TERMINAL:
        if time.time() > deadline:
            raise RuntimeError("timed out waiting for prediction")
        time.sleep(2)
        pred = _req(pred["urls"]["get"], token)
    return pred


def download(url: str, dest: Path, retries: int = 3) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    for attempt in range(retries):
        try:
            with urllib.request.urlopen(url, timeout=120) as resp:
                dest.write_bytes(resp.read())
            return
        except (urllib.error.URLError, urllib.error.HTTPError) as e:
            if attempt < retries - 1:
                time.sleep(2 ** attempt * 2)
                continue
            raise RuntimeError(f"failed to download {url}: {e}") from e


def build_input(entry: dict, args, seed: int) -> dict:
    inp = {
        "prompt": entry["prompt"],
        "output_format": args.output_format,
        "safety_tolerance": args.safety_tolerance,
        "prompt_upsampling": args.prompt_upsampling,
    }
    if args.width and args.height:
        inp["aspect_ratio"] = "custom"
        inp["width"] = args.width
        inp["height"] = args.height
    else:
        inp["aspect_ratio"] = args.aspect
    if seed is not None:
        inp["seed"] = seed
    return inp


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--prompts", required=True, type=Path)
    ap.add_argument("--issue", help="issue id; default --out is content/posts/<issue>")
    ap.add_argument("--out", type=Path, help="output root (overrides --issue)")
    ap.add_argument("--model", default="black-forest-labs/flux-1.1-pro")
    ap.add_argument("--aspect", default="2:3",
                    help="aspect_ratio; ignored when --width and --height are both set")
    ap.add_argument("--width", type=int, default=960,
                    help="custom width (multiple of 32, <=1440); 0 to use --aspect instead")
    ap.add_argument("--height", type=int, default=1440,
                    help="custom height (multiple of 32, <=1440); 0 to use --aspect instead")
    ap.add_argument("--output-format", default="png")
    ap.add_argument("--safety-tolerance", type=int, default=6)
    ap.add_argument("--prompt-upsampling", default="true",
                    type=lambda s: s.lower() not in ("false", "0", "no"))
    ap.add_argument("--seed-base", type=int, default=20260829,
                    help="per-image seed = seed_base + index; pass -1 to disable")
    ap.add_argument("--concurrency", type=int, default=3)
    ap.add_argument("--limit", type=int)
    ap.add_argument("--only", help="comma-separated slugs to (re)generate")
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    if args.out:
        out_root = args.out
    elif args.issue:
        out_root = SITE / "content" / "posts" / args.issue
    else:
        ap.error("pass --out or --issue")

    entries = json.loads(args.prompts.read_text(encoding="utf-8"))
    entries.sort(key=lambda e: e["slug"])
    if args.only:
        want = {s.strip() for s in args.only.split(",")}
        entries = [e for e in entries if e["slug"] in want]
    if args.limit:
        entries = entries[: args.limit]

    token = None if args.dry_run else load_token()
    manifest_path = args.prompts.with_suffix(".generated.json")
    manifest = []
    if manifest_path.is_file():
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            manifest = []

    tasks = []
    for i, entry in enumerate(entries):
        slug = entry["slug"]
        dest = out_root / slug / f"{slug}.png"
        seed = None if args.seed_base < 0 else args.seed_base + i
        if dest.exists() and not args.force:
            tasks.append((entry, dest, seed, "skip"))
        else:
            tasks.append((entry, dest, seed, "gen"))

    to_gen = [t for t in tasks if t[3] == "gen"]
    skipped = len(tasks) - len(to_gen)
    print(f"{len(entries)} entries · {len(to_gen)} to generate · {skipped} already present")

    if args.dry_run:
        for entry, dest, seed, _ in to_gen:
            print(f"\n--- {entry['slug']} -> {dest}")
            print(json.dumps(build_input(entry, args, seed), indent=2, ensure_ascii=False))
        return 0

    ok, failed = [], []

    def run(entry, dest, seed):
        inp = build_input(entry, args, seed)
        pred = create_prediction(args.model, inp, token)
        pred = wait_for(pred, token)
        if pred["status"] != "succeeded":
            raise RuntimeError(pred.get("error") or pred["status"])
        output = pred["output"]
        if isinstance(output, list):
            output = output[0]
        download(output, dest)
        return {
            "slug": entry["slug"],
            "prediction_id": pred.get("id"),
            "version": pred.get("version"),
            "seed": seed,
            "model": args.model,
            "output_url": output,
            "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }

    with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        futs = {pool.submit(run, e, d, s): e["slug"] for e, d, s, _ in to_gen}
        for fut in as_completed(futs):
            slug = futs[fut]
            try:
                rec = fut.result()
                ok.append(slug)
                manifest = [m for m in manifest if m.get("slug") != slug] + [rec]
                print(f"  ok    {slug}")
            except Exception as e:  # noqa: BLE001 — report and continue
                failed.append(slug)
                print(f"  FAIL  {slug}: {e}")

    manifest.sort(key=lambda m: m["slug"])
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    print(f"\n{len(ok)} ok / {len(failed)} failed / {skipped} skipped")
    if failed:
        print("failed:", ", ".join(sorted(failed)))
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())