import base64
import os
from pathlib import Path
from urllib.parse import urlparse

import requests


def save_image(item, output_dir):
    if item.get("b64_json"):
        data = base64.b64decode(item["b64_json"], validate=True)
    elif item.get("url"):
        if urlparse(item["url"]).scheme != "https":
            raise ValueError("Expected an HTTPS image URL")
        response = requests.get(item["url"], timeout=(10, 120))
        response.raise_for_status()
        data = response.content
    else:
        raise ValueError("No b64_json or url in image result")
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        suffix = ".png"
    elif data.startswith(b"\xff\xd8\xff"):
        suffix = ".jpg"
    elif data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        suffix = ".webp"
    else:
        raise ValueError("Unexpected image format; inspect the response")
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    target = output_dir / ("quickrouter-result" + suffix)
    target.write_bytes(data)
    return target


def main():
    response = requests.post(
        "https://api.quickrouter.ai/v1/images/generations",
        headers={"Authorization": "Bearer " + os.environ["QUICKROUTER_API_KEY"]},
        json={
            "model": "gpt-image-2",
            "prompt": "A small friendly blue whale on a clean white background, flat illustration, no text.",
            "n": 1,
            "size": "1024x1024",
            "quality": "low",
        },
        timeout=(10, 150),
    )
    response.raise_for_status()
    result = response.json()
    if not result.get("data"):
        raise ValueError("Empty image data")
    print(save_image(result["data"][0], "output"))
    print("usage:", result.get("usage"))


if __name__ == "__main__":
    main()