#!/usr/bin/env python3
"""
gen_face_ai.py  –  Generate assets/face_wire.png using Gemini image generation
────────────────────────────────────────────────────────────────────────────────
Uses the Gemini / Google Gen AI SDK with a native image-generation model
(gemini-2.5-flash-image or similar) to produce a wireframe face PNG ready for
display.py to load.

    pip install google-genai pillow
    python3 gen_face_ai.py

Or set the key in the environment:
    export GEMINI_API_KEY=YOUR_KEY
    python3 gen_face_ai.py
"""

import os
import sys
import io
import argparse

# ── Paste your Gemini API key here ────────────────────────────────────────────
API_KEY = 'AQ.Ab8RN6LuGwkGiKPa61jsLAEYEpJp1Yl2EkZuBWTbN9AMKxgTSw'
# ─────────────────────────────────────────────────────────────────────────────

# ── CLI ───────────────────────────────────────────────────────────────────────
ap = argparse.ArgumentParser()
ap.add_argument('--key',  default='', help='Override the hardcoded API key')
ap.add_argument('--out',  default='assets/face_wire.png')
ap.add_argument('--model', default='gemini-2.5-flash-image',
                help='Gemini image model to use (default: gemini-2.5-flash-image)')
ap.add_argument('--list-models', action='store_true',
                help='Print models that support generateContent then exit')
args = ap.parse_args()

api_key = args.key or API_KEY or os.environ.get('GEMINI_API_KEY', '')
if not api_key:
    sys.exit('ERROR: paste your key into API_KEY at the top of this file')

# ── Install check ─────────────────────────────────────────────────────────────
try:
    from google import genai
    from google.genai import types
except ImportError:
    sys.exit('Run:  pip install google-genai   then try again.')

try:
    from PIL import Image
except ImportError:
    sys.exit('Run:  pip install pillow   then try again.')

# ── Connect ───────────────────────────────────────────────────────────────────
print('Connecting to Google GenAI …')
client = genai.Client(api_key=api_key)

if args.list_models:
    print('\nModels with generateContent support:')
    for m in client.models.list():
        methods = getattr(m, 'supported_actions', None) or getattr(m, 'supported_methods', None) or []
        if 'generateContent' in methods:
            print(f'  {m.name}')
    sys.exit(0)

# ── Prompt ────────────────────────────────────────────────────────────────────
PROMPT = (
    "3D wireframe polygon mesh of a human head and face, viewed from slightly "
    "below, front-facing, neutral expression, pure black background, thin "
    "light gray lines only forming a grid over the head and face, no colour "
    "fill, no skin texture, no neck, symmetrical, sci-fi holographic display "
    "style, high contrast monochrome"
)

# ── Generate ──────────────────────────────────────────────────────────────────
print(f'Generating with {args.model} …')
response = client.models.generate_content(
    model   = args.model,
    contents= PROMPT,
    config  = types.GenerateContentConfig(
        response_modalities = ['IMAGE', 'TEXT'],
    ),
)

# ── Extract image bytes ───────────────────────────────────────────────────────
img_bytes = None
for part in response.candidates[0].content.parts:
    if part.inline_data and part.inline_data.mime_type.startswith('image/'):
        img_bytes = part.inline_data.data
        break

if img_bytes is None:
    # Print any text the model returned to help debug
    for part in response.candidates[0].content.parts:
        if hasattr(part, 'text') and part.text:
            print('Model text response:', part.text[:400])
    sys.exit('No image in response – try a different --model')

# ── Save as PNG ───────────────────────────────────────────────────────────────
os.makedirs(os.path.dirname(args.out) if os.path.dirname(args.out) else '.',
            exist_ok=True)

# Convert whatever format came back to a proper PNG
img = Image.open(io.BytesIO(img_bytes))
img.save(args.out, 'PNG')

print(f'Saved  {args.out}  ({img.width}×{img.height} px)')
print('Run  python3 display.py  to see it.')