"""robot.py — make a picture with the NVIDIA image API (FLUX.1-schnell).

Setup (one time):
    pip install requests
    Get a FREE key at https://build.nvidia.com  (top-right "Get API Key")
    Save it as the NVIDIA_API_KEY environment variable:
        export NVIDIA_API_KEY="nvapi-your_real_key_here"   # Mac / Linux
        $env:NVIDIA_API_KEY = "nvapi-your_real_key_here"   # Windows PowerShell

Run it:
    python robot.py

Companion lesson: https://krueng.ai/nvidia_api.html
"""
import base64
import os
import requests

# 1. The ADDRESS (where). The KEY (who) is read from NVIDIA_API_KEY.
API_URL = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.1-schnell"
headers = {
    "Authorization": f"Bearer {os.environ['NVIDIA_API_KEY']}",
    "Accept": "application/json",
}

# 2. Send a POST with the ORDER (body) — our picture idea.
#    cfg_scale=0 and steps=4 are required for the "schnell" (fast) FLUX variant.
response = requests.post(API_URL, headers=headers, json={
    "prompt": "a friendly robot painting a colorful picture",
    "mode": "base",
    "cfg_scale": 0,
    "width": 1024,
    "height": 1024,
    "seed": 0,
    "steps": 4,
})

# 3. If NVIDIA refused the request, SHOW the reason and stop —
#    don't pretend success when the response is an error JSON.
if not response.ok:
    print(f"NVIDIA rejected the request ({response.status_code}):")
    print(response.text)
    raise SystemExit(1)

# 4. NVIDIA hides the picture inside a JSON envelope as a base64 string.
#    Decode it and save the bytes to a normal .jpg file.
data = response.json()
image_bytes = base64.b64decode(data["artifacts"][0]["base64"])
with open("robot.jpg", "wb") as file:
    file.write(image_bytes)

print("Done! Open robot.jpg to see your picture.")