#!/usr/bin/env python3
"""Two-step approach: 1) Generate empty MacBook low-light scene, 2) Edit screen content in."""

import json
import base64
import sys
import urllib.request
import urllib.error
from pathlib import Path

API_KEY = "AIzaSyDXqYZInk83iVV4mD29pSuHQKbgkiI1x9Q"
MODEL = "gemini-2.5-flash-image"

STEP = sys.argv[1] if len(sys.argv) > 1 else "1"

def load_image(path):
    data = Path(path).read_bytes()
    mime = "image/png" if path.endswith(".png") else "image/jpeg"
    return base64.b64encode(data).decode("utf-8"), mime

def generate(prompt, output, input_images=None):
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:generateContent?key={API_KEY}"
    parts = [{"text": prompt}]
    if input_images:
        for img_path in input_images:
            b64, mime = load_image(img_path)
            parts.append({"inlineData": {"mimeType": mime, "data": b64}})
    
    payload = {
        "contents": [{"parts": parts}],
        "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]}
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    
    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            result = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8")
        print(f"HTTP {e.code}: {body[:500]}")
        return False
    
    for candidate in result.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            if "inlineData" in part:
                img_data = base64.b64decode(part["inlineData"]["data"])
                with open(output, "wb") as f:
                    f.write(img_data)
                print(f"Saved: {output} ({len(img_data)} bytes)")
                return True
            elif "text" in part:
                print(f"Text: {part['text'][:200]}")
    
    print("No image generated")
    print(json.dumps(result, indent=2)[:800])
    return False

if STEP == "1":
    print("Step 1: Generating empty MacBook low-light scene...")
    prompt = """Create a photorealistic photograph of a MacBook Pro laptop open on a dark desk in a very low-light environment. The laptop screen shows a bright white webpage with a clean grid layout of colorful cards and text content — a modern SaaS dashboard interface. The bright screen is the only light source, casting a soft warm glow onto the dark matte desk surface. The laptop keyboard is dimly backlit. Shot from above at 30 degrees, slightly off-center left. 16:9 composition. Ultra-premium product photography quality. Near pitch-black background fading to pure dark."""
    generate(prompt, "/root/.openclaw/workspace/work/pitches/brandwatch/visuals/lowlight-base.png")

elif STEP == "2":
    print("Step 2: Editing screen content into the mockup...")
    prompt = """Edit this laptop mockup image: replace the content on the laptop screen with the content from the second image. Keep the screen exactly the same size and perspective. Preserve all the lighting, shadows, and glow effects from the original mockup. The screen should display the social media dashboard content clearly and sharply."""
    generate(
        prompt,
        "/root/.openclaw/workspace/work/pitches/brandwatch/visuals/production-lowlight.png",
        [
            "/root/.openclaw/workspace/work/pitches/brandwatch/visuals/lowlight-base.png",
            "/root/.openclaw/workspace/work/pitches/brandwatch/visuals/production-assets-screen.png"
        ]
    )