#!/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?

    The question is ONLY "has the daily allocation run out". Anything else the
    API says means the account is alive and spending is possible, so it counts as
    budget available — including a rejected prompt. That distinction is not
    hypothetical: Cloudflare's safety filter rejected the literal prompt "a grey
    square" as NSFW (HTTP 400, code 3030), and because the old probe let every
    non-429 fall through to SystemExit, a perfectly healthy account with a full
    10,000 neurons reported as unusable.
    """
    try:
        _post("a plain wooden chair, product photo", 4, 1)
        return True
    except Exhausted:
        return False
    except RateLimited:
        # Burst-limited means the account is alive and has budget.
        return True
    except SystemExit as exc:
        # A 4xx that is not the daily cap: the account works, this prompt did not.
        print(f"cloudflare reachable, prompt refused: {exc}", file=sys.stderr)
        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()