#!/usr/bin/env python3
"""
Open Flow in headed browser for the user to interact with.
Captures all aisandbox API requests using page.on("request"/"response").
This approach persists data even when browser closes.
"""
import json
import os
import sys
import time
from dotenv import load_dotenv
load_dotenv()
from playwright.sync_api import sync_playwright

s = os.getenv("GOOGLE_FLOW_SESSION_TOKEN", "")
c = os.getenv("GOOGLE_FLOW_CSRF_TOKEN", "")

if not s or not c:
    print("Missing GOOGLE_FLOW_SESSION_TOKEN or GOOGLE_FLOW_CSRF_TOKEN in .env")
    sys.exit(1)

# Store captured data in Python (not in browser JS)
captured_requests = []
pending_requests = {}

def handle_request(request):
    url = request.url
    if "aisandbox" in url:
        body = request.post_data
        parsed = None
        if body:
            try:
                parsed = json.loads(body)
            except:
                pass
        pending_requests[url] = {
            "url": url,
            "method": request.method,
            "requestBody": parsed,
            "timestamp": time.time(),
        }

def handle_response(response):
    url = response.url
    if url in pending_requests:
        req_data = pending_requests.pop(url)
        try:
            resp_body = response.json()
        except:
            resp_body = None
        req_data["status"] = response.status
        req_data["responseBody"] = resp_body
        captured_requests.append(req_data)
        
        # Save immediately after each capture
        with open("captured_i2v_requests.json", "w") as f:
            json.dump(captured_requests, f, indent=2)
        
        if "video" in url.lower():
            print(f"  [Captured: {url.split('/')[-1][:50]}]")

with sync_playwright() as p:
    browser = p.chromium.launch(headless=False)
    context = browser.new_context(
        viewport={"width": 1400, "height": 900},
        locale="pt-BR",
    )

    context.add_cookies([
        {
            "name": "__Secure-next-auth.session-token",
            "value": s,
            "domain": "labs.google",
            "path": "/",
            "secure": True,
            "httpOnly": True,
            "sameSite": "Lax",
        },
        {
            "name": "__Host-next-auth.csrf-token",
            "value": c,
            "domain": "labs.google",
            "path": "/",
            "secure": True,
            "httpOnly": True,
            "sameSite": "Lax",
        },
    ])

    page = context.new_page()
    
    # Use Playwright's built-in request/response interception
    page.on("request", handle_request)
    page.on("response", handle_response)

    print("Opening Flow video tool...")
    print(">>> Configure: Veo 3.1 Fast, 1 response, portrait 9:16")
    print(">>> Add an image as start frame, type prompt, click generate")
    print(">>> After generation starts, close the browser window")
    print(">>> Data is saved automatically as you interact!")
    print()

    page.goto("https://labs.google/fx/tools/video-fx", wait_until="networkidle", timeout=60000)

    # Wait for user to close browser
    try:
        while True:
            if page.is_closed():
                break
            time.sleep(1)
    except Exception:
        pass

    try:
        browser.close()
    except:
        pass

print(f"\n{'='*60}")
print(f"Captured {len(captured_requests)} total aisandbox requests")

# Filter video requests
video_requests = [r for r in captured_requests 
                  if any(kw in r.get("url", "") for kw in ["video", "Video", "batchAsync"])]

for req in video_requests:
    print(f"\n{'='*60}")
    print(f"URL: {req.get('url')}")
    print(f"Method: {req.get('method')} | Status: {req.get('status')}")
    rb = req.get('requestBody')
    if rb:
        print(f"Request Body:\n{json.dumps(rb, indent=2)[:2000]}")

if video_requests:
    with open("i2v_request_body.json", "w") as f:
        json.dump(video_requests, f, indent=2)
    print(f"\nVideo requests saved to i2v_request_body.json")
else:
    print("\nNo video requests captured.")