#!/usr/bin/env python3
"""Generate images using Gemini with Keith's reference photos."""
import base64, json, sys, os, requests

API_KEY = "AIzaSyCX7Jwk-iHYdIjw5qsAb94bZrScWJw84q8"
MODEL = "nano-banana-pro-preview"
PHOTO_DIR = os.path.dirname(os.path.abspath(__file__))
OUTPUT = os.path.join(PHOTO_DIR, "output.png")

def load_image_b64(path):
    with open(path, "rb") as f:
        return base64.standard_b64encode(f.read()).decode()

def generate(prompt, photo_paths, output_path=OUTPUT):
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={API_KEY}"
    
    parts = []
    # Add reference photos
    for p in photo_paths:
        parts.append({
            "inline_data": {
                "mime_type": "image/jpeg",
                "data": load_image_b64(p)
            }
        })
    # Add text prompt
    parts.append({"text": prompt})
    
    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "temperature": 1.0
        }
    }
    
    resp = requests.post(url, json=payload, timeout=120)
    if resp.status_code != 200:
        print(f"Error {resp.status_code}: {resp.text[:500]}")
        sys.exit(1)
    
    data = resp.json()
    # Extract image from response
    for candidate in data.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            if "inlineData" in part:
                img_data = base64.standard_b64decode(part["inlineData"]["data"])
                with open(output_path, "wb") as f:
                    f.write(img_data)
                print(f"Image saved to {output_path}")
                return
            if "text" in part:
                print(f"Text: {part['text']}")
    
    print("No image in response")
    print(json.dumps(data, indent=2)[:1000])

if __name__ == "__main__":
    prompt = sys.argv[1] if len(sys.argv) > 1 else "Generate a portrait of this person as an astronaut on Mars"
    photos = [
        os.path.join(PHOTO_DIR, "headshot-full.jpg"),
        os.path.join(PHOTO_DIR, "headshot-close.jpg"),
    ]
    generate(prompt, photos)