#!/usr/bin/env python3
"""End-to-end test: cookies → session token → Bearer auth → generate image."""
import base64
import json
import os
import sys
import time

from dotenv import load_dotenv
load_dotenv()

import httpx

s = os.getenv("GOOGLE_FLOW_SESSION_TOKEN", "")
c = os.getenv("GOOGLE_FLOW_CSRF_TOKEN", "")
cookie_str = f"__Secure-next-auth.session-token={s}; __Host-next-auth.csrf-token={c}"

print(f"Session: {s[:40]}...")
print(f"CSRF: {c[:40]}...")

client = httpx.Client(timeout=60, follow_redirects=True)
headers = {
    "Origin": "https://labs.google",
    "Referer": "https://labs.google/fx/tools/image-fx",
    "Content-Type": "application/json",
    "Cookie": cookie_str,
}

# Step 1: Get access token from session
print("\n--- Step 1: Get access token ---")
r = client.get("https://labs.google/fx/api/auth/session", headers=headers)
print(f"Status: {r.status_code}")
data = r.json()
print(f"Keys: {list(data.keys())}")
access_token = data.get("access_token") or data.get("accessToken")
if access_token:
    print(f"Access token: {access_token[:50]}...")
    print(f"Expires: {data.get('expires')}")
    print(f"User: {data.get('user', {}).get('email', 'N/A')}")
else:
    print(f"NO ACCESS TOKEN! Full response: {json.dumps(data)[:500]}")
    sys.exit(1)

# Step 2: Generate image
print("\n--- Step 2: Generate image ---")
auth_headers = {**headers, "Authorization": f"Bearer {access_token}"}
body = {
    "userInput": {
        "candidatesCount": 1,
        "prompts": ["A cute cat sitting on a windowsill, photorealistic, 8K"],
        "seed": 42,
    },
    "clientContext": {
        "sessionId": f";{int(time.time()*1000)}",
        "tool": "IMAGE_FX",
    },
    "modelInput": {
        "modelNameType": "IMAGEN_3_5",
    },
    "aspectRatio": "IMAGE_ASPECT_RATIO_LANDSCAPE",
}

r = client.post("https://aisandbox-pa.googleapis.com/v1:runImageFx",
    json=body, headers=auth_headers)
print(f"Status: {r.status_code}")
print(f"Content-Type: {r.headers.get('content-type')}")

if r.status_code == 200:
    result = r.json()
    print(f"Response keys: {list(result.keys())}")
    panels = result.get("imagePanels", [])
    print(f"Panels: {len(panels)}")
    if panels:
        images = panels[0].get("generatedImages", [])
        print(f"Images: {len(images)}")
        if images:
            encoded = images[0].get("encodedImage", "")
            if encoded:
                img_data = base64.b64decode(encoded)
                out = "output/test_google_flow.png"
                os.makedirs("output", exist_ok=True)
                with open(out, "wb") as f:
                    f.write(img_data)
                print(f"\nSaved: {out} ({len(img_data)} bytes)")
            else:
                print(f"Image keys: {list(images[0].keys())}")
else:
    print(f"Body: {r.text[:500]}")

client.close()