#!/usr/bin/env python3
"""Generate brand-style AI images for the Meridian post-banner well.

Uses the Gemini flash image model exposed through the local LiteLLM proxy
(http://litellm:4000, team key from config.yaml model.api_key).

Usage:
  generate_ai_image.py "<prompt>" <ratio> <output.png>
    ratio: 4:5 | 1:1 | 16:9 | 9:16

Brand defaults: navy #203462 accents, amber #F8A538 highlights, white dominant,
no text/logos/watermarks/people, clean negative space.
"""
import sys, os, json, base64, re, urllib.request, urllib.error

def main():
    if len(sys.argv) < 4:
        print(__doc__)
        sys.exit(2)
    prompt, ratio, out = sys.argv[1], sys.argv[2], sys.argv[3]
    if not os.path.isabs(out):
        # default destination: this bundle's generated ai-images folder
        bundle = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
        out = os.path.join(bundle, "C.Purpose", "c2.Marketing", "brandkit", "assets", "generated", "ai-images", out)
    model = os.environ.get("BRAND_IMAGE_MODEL", "Nostics_Image")
    cfg = open("/opt/data/config.yaml").read()
    m = re.search(r'^\s*api_key:\s*["\']?([^"\'\n]+)', cfg, re.M)
    key = m.group(1).strip()
    req = urllib.request.Request(
        "http://litellm:4000/v1/images/generations",
        data=json.dumps({"model": model, "prompt": prompt, "n": 1, "size": ratio}).encode(),
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        method="POST")
    try:
        with urllib.request.urlopen(req, timeout=240) as r:
            data = json.loads(r.read())
    except urllib.error.HTTPError as e:
        print(f"ERROR {e.code} on model '{model}': {e.read().decode()[:800]}", file=sys.stderr)
        print("Check the alias against GET http://litellm:4000/v1/models, then override with "
              "BRAND_IMAGE_MODEL=<alias>. Known history: 'gemini-3.1-flash-image' 400s "
              "(renamed 2026-10-15) and 'Nostics_Image' has refused with "
              "'LLM Provider NOT provided' — if no alias works, reuse an existing image from "
              "assets/generated/ai-images/ rather than stalling the post.", file=sys.stderr)
        sys.exit(1)
    items = data.get("data", [])
    if not items:
        print("ERROR: no image data in response", file=sys.stderr)
        sys.exit(1)
    it = items[0]
    os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True)
    if it.get("b64_json"):
        with open(out, "wb") as f:
            f.write(base64.b64decode(it["b64_json"]))
        print(f"saved {out}")
    elif it.get("url"):
        with urllib.request.urlopen(it["url"], timeout=120) as r, open(out, "wb") as f:
            f.write(r.read())
        print(f"saved {out} (from url)")
    else:
        print("ERROR: no b64_json or url", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()