#!/usr/bin/env python3
"""cf_flux.py — FLUX-1-schnell images from Cloudflare Workers AI.

    python3 tools/gen/cf_flux.py <out.png> "<prompt>" [--steps 6] [--seed 42]
    python3 tools/gen/cf_flux.py --probe          # is there budget left today?

Workers AI gives 10,000 Neurons free every day (resets 00:00 UTC), which is the
cheapest image generation on the fleet — so it goes FIRST, ahead of the farm.
When it runs dry, callers fall back to MODELBEAST `flux_local`, which is slower
but genuinely unlimited.

The whole reason this exists rather than the one-liner in toastsim: urllib
raises HTTPError on 4xx and throws the response BODY away, so a 429 is
indistinguishable between "you have burned today's neurons, go to the farm" and
"you sent three at once, wait a second". Those want opposite responses, and
guessing wrong either wastes the free tier or wedges the batch. Here the body is
read and classified.
"""
from __future__ import annotations

import base64
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import creds  # noqa: E402

MODEL = "@cf/black-forest-labs/flux-1-schnell"
# schnell is a distilled 4-step model; past ~8 steps you pay neurons for nothing.
DEFAULT_STEPS = 6


class Exhausted(RuntimeError):
    """Today's free neurons are gone. Fall back to the farm; retrying won't help."""


class RateLimited(RuntimeError):
    """Too fast, not too much. Backing off DOES help."""


def _post(prompt: str, steps: int, seed: int) -> bytes:
    account = creds.get("CLOUDFLARE_ACCOUNT_ID")
    token = creds.get("CLOUDFLARE_API_TOKEN")
    req = urllib.request.Request(
        f"https://api.cloudflare.com/client/v4/accounts/{account}/ai/run/{MODEL}",
        data=json.dumps({"prompt": prompt, "steps": steps, "seed": seed}).encode(),
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
    )
    try:
        raw = urllib.request.urlopen(req, timeout=180).read()
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="ignore")[:600]
        low = body.lower()
        if exc.code == 429:
            # Cloudflare says "capacity temporarily exceeded" for burst limits and
            # names the neuron/daily allocation when the free tier is spent.
            if "neuron" in low or "daily" in low or "limit" in low and "rate" not in low:
                raise Exhausted(body) from None
            raise RateLimited(body) from None
        if exc.code in (401, 403):
            raise SystemExit(f"Cloudflare rejected the token (HTTP {exc.code}): {body}")
        raise SystemExit(f"Cloudflare HTTP {exc.code}: {body}") from None

    payload = json.loads(raw)
    if not payload.get("success"):
        errors = json.dumps(payload.get("errors"))
        if "neuron" in errors.lower() or "limit" in errors.lower():
            raise Exhausted(errors)
        raise SystemExit("Cloudflare error: " + errors)
    return base64.b64decode(payload["result"]["image"])


def generate(out: Path, prompt: str, steps: int = DEFAULT_STEPS, seed: int = 0,
             *, attempts: int = 4) -> Path:
    """One image to `out`. Retries burst limits, gives up fast on exhaustion."""
    delay = 2.0
    for attempt in range(1, attempts + 1):
        try:
            out.parent.mkdir(parents=True, exist_ok=True)
            out.write_bytes(_post(prompt, steps, seed))
            return out
        except RateLimited:
            if attempt == attempts:
                raise
            time.sleep(delay)
            delay *= 2
    raise RateLimited("exhausted retries")


def probe() -> bool:
    """Cheapest possible call — is there budget left right now?"""
    try:
        _post("a grey square", 1, 1)
        return True
    except Exhausted:
        return False
    except RateLimited:
        # Burst-limited means the account is alive and has budget.
        return True


def main() -> None:
    if "--probe" in sys.argv:
        ok = probe()
        print("cloudflare flux: BUDGET AVAILABLE" if ok else "cloudflare flux: EXHAUSTED for today")
        sys.exit(0 if ok else 3)
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    out, prompt = Path(sys.argv[1]), sys.argv[2]
    steps = int(sys.argv[sys.argv.index("--steps") + 1]) if "--steps" in sys.argv else DEFAULT_STEPS
    seed = int(sys.argv[sys.argv.index("--seed") + 1]) if "--seed" in sys.argv else 0
    try:
        print(generate(out, prompt, steps, seed))
    except Exhausted as exc:
        sys.exit(f"cloudflare free tier exhausted for today — use MODELBEAST flux_local.\n{exc}")


if __name__ == "__main__":
    main()