#!/usr/bin/env python3
"""
Fully automate I2V video generation on Flow:
1. Navigate with cookies
2. Create new project
3. Configure Veo 3.1 Fast, 1 response, portrait 9:16
4. Upload a test image as start frame
5. Type prompt and click generate
6. Capture the request body
"""
import json
import os
import sys
import time
from pathlib import Path
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 cookies in .env")
    sys.exit(1)

# Store captured data
captured_requests = []

def handle_request(request):
    url = request.url
    # Capture ALL aisandbox requests to find the image upload endpoint
    if "aisandbox" in url or "upload" in url.lower() or "media" in url.lower():
        body = request.post_data
        parsed = None
        if body:
            try:
                parsed = json.loads(body)
            except:
                # Might be binary/form data
                parsed = f"<binary data: {len(body)} bytes>"
        captured_requests.append({
            "url": url,
            "method": request.method,
            "requestBody": parsed,
            "content_type": request.headers.get("content-type", ""),
        })
        print(f"  [CAPTURED] {url.split('/')[-1][:60]}")
        # Save immediately
        with open("all_requests.json", "w") as f:
            json.dump(captured_requests, f, indent=2, default=str)

# Find a test image to upload
test_images = list(Path(".").glob("*.png")) + list(Path(".").glob("*.jpg")) + list(Path("output").glob("*.png"))
if not test_images:
    # Create a simple test image
    print("Creating test image...")
    from PIL import Image
    img = Image.new('RGB', (1024, 1024), color=(100, 150, 200))
    img.save("test_frame.png")
    test_image = Path("test_frame.png").absolute()
else:
    test_image = test_images[0].absolute()
print(f"Using test image: {test_image}")

with sync_playwright() as p:
    browser = p.chromium.launch(headless=False)  # Visible for debugging
    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()
    page.on("request", handle_request)

    print("Navigating to Flow...")
    page.goto("https://labs.google/fx/tools/video-fx", wait_until="networkidle", timeout=60000)
    page.wait_for_timeout(3000)
    
    # Take screenshot to see current state
    page.screenshot(path="flow_step1_landing.png")
    print("Screenshot: flow_step1_landing.png")

    # Click "Create with Flow" or "Novo projeto"
    try:
        page.locator("button:has-text('Create with Flow'), button:has-text('Criar com Flow')").first.click(timeout=5000)
        print("Clicked 'Create with Flow'")
        page.wait_for_timeout(5000)
    except:
        print("No 'Create with Flow' button, trying 'New project'...")
        try:
            page.locator("button:has-text('Novo projeto'), button:has-text('New project')").first.click(timeout=5000)
            print("Clicked 'New project'")
            page.wait_for_timeout(5000)
        except:
            print("No project button found")
    
    page.screenshot(path="flow_step2_after_click.png")
    print("Screenshot: flow_step2_after_click.png")

    # Wait for the editor to load - look for textarea
    print("Waiting for editor...")
    try:
        page.locator("textarea").first.wait_for(state="visible", timeout=30000)
        print("Editor loaded!")
    except:
        print("Textarea not found, checking page state...")
        page.screenshot(path="flow_step3_no_textarea.png")
        
        # Maybe we need to click something else
        # Look for any clickable element that might open the editor
        buttons = page.locator("button").all()
        print(f"Found {len(buttons)} buttons:")
        for i, btn in enumerate(buttons[:10]):
            try:
                txt = btn.inner_text()[:50]
                print(f"  {i}: {txt}")
            except:
                pass

    # Now configure the settings
    # Look for settings/config dropdowns
    print("\nLooking for settings dropdowns...")
    
    # Find "Respostas por comando" dropdown and set to 1
    try:
        # Look for dropdown with "2" (default) and change to "1"
        page.locator("text=Respostas por comando").first.wait_for(timeout=5000)
        # Click the dropdown near it
        page.locator("text=Respostas por comando").locator("..").locator("button, [role='combobox']").first.click()
        page.wait_for_timeout(500)
        page.locator("text='1'").first.click()
        print("Set responses to 1")
    except Exception as e:
        print(f"Could not set responses: {e}")

    # Look for model selector
    try:
        page.locator("text=Modelo").first.wait_for(timeout=3000)
        page.locator("text=Modelo").locator("..").locator("button, [role='combobox']").first.click()
        page.wait_for_timeout(500)
        # Select Veo 3.1 Fast
        page.locator("text='Veo 3.1 - Fast'").first.click()
        print("Set model to Veo 3.1 Fast")
    except Exception as e:
        print(f"Could not set model: {e}")

    # Look for aspect ratio
    try:
        page.locator("text=Proporção").first.wait_for(timeout=3000)
        page.locator("text=Proporção").locator("..").locator("button, [role='combobox']").first.click()
        page.wait_for_timeout(500)
        page.locator("text='Retrato'").first.click()
        print("Set aspect ratio to Portrait")
    except Exception as e:
        print(f"Could not set aspect ratio: {e}")

    page.screenshot(path="flow_step4_settings.png")

    # Upload image as start frame
    print("\nLooking for image upload...")
    try:
        # Look for file input or upload button
        file_input = page.locator("input[type='file']").first
        if file_input:
            file_input.set_input_files(str(test_image))
            print(f"Uploaded image: {test_image}")
            page.wait_for_timeout(3000)
    except Exception as e:
        print(f"Could not upload image: {e}")
        # Try clicking an upload button
        try:
            page.locator("button:has-text('upload'), button:has-text('Carregar'), [aria-label*='upload']").first.click()
            page.wait_for_timeout(1000)
            page.locator("input[type='file']").first.set_input_files(str(test_image))
            print("Uploaded via button click")
        except:
            print("No upload mechanism found")

    page.screenshot(path="flow_step5_after_upload.png")

    # Type prompt
    print("\nTyping prompt...")
    try:
        textarea = page.locator("textarea").first
        textarea.click()
        page.wait_for_timeout(300)
        textarea.fill("A cute orange cat walking slowly, cinematic lighting, 4K quality")
        print("Typed prompt")
    except Exception as e:
        print(f"Could not type prompt: {e}")

    page.screenshot(path="flow_step6_prompt.png")

    # Click generate
    print("\nClicking generate...")
    try:
        page.locator("button:has-text('Criar'), button:has-text('Create'), button:has-text('Generate')").first.click(timeout=5000)
        print("Clicked generate!")
    except Exception as e:
        print(f"Could not click generate: {e}")

    # Wait for the request to be captured
    print("\nWaiting for video generation request...")
    for i in range(30):
        page.wait_for_timeout(2000)
        if captured_requests:
            print(f"Captured {len(captured_requests)} video request(s)!")
            break
        print(f"  Waiting... ({i*2}s)")

    page.screenshot(path="flow_step7_generating.png")

    # Keep browser open for a bit to capture more
    page.wait_for_timeout(5000)

    browser.close()

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

if captured_requests:
    print("\nI2V Request Body:")
    print(json.dumps(captured_requests[0].get("requestBody"), indent=2))
    print(f"\nSaved to i2v_request_body.json")
else:
    print("No video requests captured. Check the screenshots for debugging.")