"""
FlowBrowserEngine — Headless Playwright engine for Google Labs Flow.

Reverse-engineered API (Feb 2026):
  Endpoint: POST aisandbox-pa.googleapis.com/v1/projects/{project_id}/flowMedia:batchGenerateImages
  Model:    GEM_PIX_2 (displayed as "Nano Banana Pro" in UI)
  Auth:     Bearer token (from session endpoint) + reCAPTCHA (browser-generated)
  Response: Signed GCS URLs for each generated image

Strategy:
1. Keep headless Chromium alive on a Flow project page
2. Install JS fetch() interceptor to capture reCAPTCHA + auth from real requests
3. For each generation: fill prompt → click submit → wait for API response
4. Download images from the returned signed URLs
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import random
import time
from pathlib import Path
from typing import Optional

import httpx

logger = logging.getLogger(__name__)

FLOW_GALLERY = "https://labs.google/fx/pt/tools/flow"  # Main Flow page (not gallery)
ASPECT_RATIOS = {
    "portrait": "IMAGE_ASPECT_RATIO_PORTRAIT",
    "landscape": "IMAGE_ASPECT_RATIO_LANDSCAPE",
    "square": "IMAGE_ASPECT_RATIO_SQUARE",
}

# ── Stealth patches to evade reCAPTCHA Enterprise bot detection ───────────────
_STEALTH_JS = """
// 1. Hide navigator.webdriver
Object.defineProperty(navigator, 'webdriver', { get: () => undefined });

// 2. Override Permissions API to look like real Chrome
const originalQuery = window.navigator.permissions.query;
window.navigator.permissions.query = (parameters) => (
    parameters.name === 'notifications'
        ? Promise.resolve({ state: Notification.permission })
        : originalQuery(parameters)
);

// 3. Fake plugins (Chrome always has at least PDF plugins)
Object.defineProperty(navigator, 'plugins', {
    get: () => {
        const plugins = [
            { name: 'Chrome PDF Plugin', filename: 'internal-pdf-viewer', description: 'Portable Document Format' },
            { name: 'Chrome PDF Viewer', filename: 'mhjfbmdgcfjbbpaeojofohoefgiehjai', description: '' },
            { name: 'Native Client', filename: 'internal-nacl-plugin', description: '' },
        ];
        plugins.length = 3;
        plugins.refresh = () => {};
        plugins.item = (i) => plugins[i] || null;
        plugins.namedItem = (n) => plugins.find(p => p.name === n) || null;
        return plugins;
    }
});

// 4. Fake languages
Object.defineProperty(navigator, 'languages', { get: () => ['pt-BR', 'pt', 'en-US', 'en'] });

// 5. Fake Chrome runtime (reCAPTCHA checks window.chrome)
if (!window.chrome) {
    window.chrome = {};
}
if (!window.chrome.runtime) {
    window.chrome.runtime = {
        connect: () => {},
        sendMessage: () => {},
        id: undefined,
    };
}

// 6. Override connection info to hide automation
Object.defineProperty(navigator, 'connection', {
    get: () => ({
        effectiveType: '4g',
        rtt: 50,
        downlink: 10,
        saveData: false,
    })
});

// 7. Fake WebGL vendor/renderer (headless often has different values)
const getParameter = WebGLRenderingContext.prototype.getParameter;
WebGLRenderingContext.prototype.getParameter = function(param) {
    if (param === 37445) return 'Google Inc. (NVIDIA)';
    if (param === 37446) return 'ANGLE (NVIDIA, NVIDIA GeForce GTX 1080 Direct3D11 vs_5_0 ps_5_0, D3D11)';
    return getParameter.call(this, param);
};

// 8. Fake platform
Object.defineProperty(navigator, 'platform', { get: () => 'Linux x86_64' });

// 9. Fake hardware concurrency
Object.defineProperty(navigator, 'hardwareConcurrency', { get: () => 8 });

// 10. Fake device memory
Object.defineProperty(navigator, 'deviceMemory', { get: () => 8 });
"""  # noqa: E501

# JS fetch interceptor — captures aisandbox POST requests,
# overrides imageAspectRatio via window.__flowAspectRatio, and
# stores the batchGenerateImages response for Python to read.
_FETCH_INTERCEPTOR = """
(() => {
    if (window.__flowInterceptorReady) return;
    window.__flowCaptured = [];
    window.__flowLastGeneration = null;
    window.__flowAspectRatio = null;  // set from Python to override aspect ratio
    window.__flowReferenceMediaId = null;  // set from Python for coverage shot reference
    window.__flowReferenceWeight = 0.85;   // reference image influence weight
    window.__flowLastMediaId = null;       // last captured mediaGenerationId
    const origFetch = window.fetch;
    window.fetch = async function(...args) {
        const [url, opts] = args;
        const urlStr = typeof url === 'string' ? url : url?.url || '';
        if (urlStr.includes('aisandbox-pa') && opts?.method?.toUpperCase() === 'POST') {
            const entry = {url: urlStr, headers: {}, body: null, ts: Date.now()};
            if (opts.headers) {
                if (opts.headers instanceof Headers)
                    opts.headers.forEach((v, k) => { entry.headers[k] = v; });
                else if (typeof opts.headers === 'object')
                    entry.headers = {...opts.headers};
            }
            // Override aspect ratio + limit to 1 image in batchGenerateImages requests
            if (urlStr.includes('batchGenerateImages') && opts.body) {
                try {
                    const body = JSON.parse(opts.body);
                    if (body.requests && body.requests.length > 0) {
                        // Keep only the first request to generate 1 image
                        body.requests = [body.requests[0]];
                        if (window.__flowAspectRatio) {
                            body.requests[0].imageAspectRatio = window.__flowAspectRatio;
                        }
                        // Inject reference image for coverage shot consistency
                        if (window.__flowReferenceMediaId) {
                            body.requests[0].imageInputs = [{
                                mediaId: window.__flowReferenceMediaId,
                                weight: window.__flowReferenceWeight || 0.85
                            }];
                        } else {
                            body.requests[0].imageInputs = [];
                        }
                        // Log the prompt being sent to API
                        const promptText = body.requests[0]?.prompt || body.requests[0]?.text || 'NO_PROMPT_FOUND';
                        console.log('[FlowInterceptor] Prompt in API request (' + promptText.length + ' chars):', promptText.substring(0, 300));
                        window.__flowLastPromptSent = promptText;
                    }
                    opts.body = JSON.stringify(body);
                } catch (e) {
                    console.error('[FlowInterceptor] Error parsing body:', e);
                }
            }
            if (opts.body) {
                try { entry.body = JSON.parse(opts.body); } catch {}
            }
            window.__flowCaptured.push(entry);
        }
        const response = await origFetch.apply(this, args);
        if (urlStr.includes('batchGenerateImages')) {
            try {
                const clone = response.clone();
                const data = await clone.json();
                window.__flowLastGeneration = {url: urlStr, response: data, ts: Date.now()};
                // Capture mediaGenerationId for coverage shot reference
                try {
                    const mediaId = data?.media?.[0]?.image?.generatedImage?.mediaGenerationId;
                    if (mediaId) window.__flowLastMediaId = mediaId;
                } catch {}
            } catch {}
        }
        return response;
    };
    window.__flowInterceptorReady = true;
})();
"""

# Maximum prompt length for Google Flow (empirically determined)
# Flow works best with shorter, focused prompts
FLOW_MAX_PROMPT_LENGTH = 1200


def _condense_prompt_for_flow(prompt: str, max_length: int = FLOW_MAX_PROMPT_LENGTH) -> str:
    """Condense a long prompt to fit Flow's optimal length.
    
    Strategy:
    1. If prompt is short enough, return as-is
    2. Otherwise, prioritize: character description > environment > camera/lighting
    3. Remove redundant phrases and compress
    """
    if len(prompt) <= max_length:
        return prompt
    
    logger.warning(
        "FlowBrowserEngine: Prompt too long (%d chars), condensing to %d",
        len(prompt), max_length
    )
    
    # Split into paragraphs/sections
    sections = prompt.split('\n\n')
    if len(sections) == 1:
        sections = prompt.split('. ')
    
    # Priority keywords to keep
    priority_keywords = [
        'skeleton', 'skull', 'transparent', 'bone', 'yellow iris',
        'wearing', 'environment', 'cyclorama', 'camera', 'lighting',
        'photorealistic'
    ]
    
    # Build condensed prompt keeping priority content
    condensed_parts = []
    current_length = 0
    
    for section in sections:
        section = section.strip()
        if not section:
            continue
            
        # Check if section contains priority keywords
        has_priority = any(kw.lower() in section.lower() for kw in priority_keywords)
        
        if has_priority or current_length < max_length // 2:
            if current_length + len(section) + 2 <= max_length:
                condensed_parts.append(section)
                current_length += len(section) + 2
            elif current_length < max_length - 100:
                # Truncate this section to fit
                remaining = max_length - current_length - 50
                if remaining > 100:
                    condensed_parts.append(section[:remaining] + "...")
                break
    
    result = '. '.join(condensed_parts) if len(sections) > 1 else '\n\n'.join(condensed_parts)
    
    # Ensure it ends with photorealistic if original had it
    if 'photorealistic' in prompt.lower() and 'photorealistic' not in result.lower():
        result = result.rstrip('.') + ". Photorealistic cinematic realism."
    
    logger.info("FlowBrowserEngine: Condensed prompt to %d chars", len(result))
    return result


class FlowBrowserEngine:
    """
    Headless browser engine for Flow image generation with Nano Banana Pro.

    Usage (async):
        engine = FlowBrowserEngine(session_token, csrf_token)
        await engine.start()
        path = await engine.generate("a cat", Path("out.png"))
        await engine.stop()
    """

    def __init__(self, session_token: str, csrf_token: str):
        self.session_token = session_token
        self.csrf_token = csrf_token
        self._browser = None
        self._context = None
        self._page = None
        self._pw = None
        self._ready = False
        self._lock = asyncio.Lock()
        self._project_id: Optional[str] = None
        self._http_client: httpx.AsyncClient | None = None
        self._last_media_id: str | None = None

    @property
    def is_running(self) -> bool:
        return self._ready and self._browser is not None

    @property
    def project_id(self) -> Optional[str]:
        return self._project_id

    @property
    def last_media_id(self) -> str | None:
        """The mediaGenerationId from the last successful image generation."""
        return self._last_media_id

    # ── Lifecycle ──────────────────────────────────────────────────────────────

    async def start(self, timeout_s: int = 60) -> bool:
        """Launch browser, navigate to Flow, enter a project in image mode."""
        if self._ready:
            return True

        from playwright.async_api import async_playwright

        try:
            self._pw = await async_playwright().start()
            self._browser = await self._pw.chromium.launch(
                headless=True,
                args=[
                    "--no-sandbox",
                    "--disable-blink-features=AutomationControlled",
                    "--disable-dev-shm-usage",
                    "--disable-infobars",
                    "--window-size=1280,900",
                    "--start-maximized",
                    "--disable-extensions",
                    "--disable-component-extensions-with-background-pages",
                    "--disable-background-networking",
                    "--disable-default-apps",
                ],
            )
            self._context = await self._browser.new_context(
                user_agent=(
                    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
                    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
                ),
                viewport={"width": 1280, "height": 900},
                locale="pt-BR",
            )
            # Inject stealth patches before any page JS runs
            await self._context.add_init_script(_STEALTH_JS)
            await self._context.add_cookies(
                [
                    {
                        "name": "__Secure-next-auth.session-token",
                        "value": self.session_token,
                        "domain": "labs.google",
                        "path": "/",
                        "secure": True,
                        "httpOnly": True,
                        "sameSite": "Lax",
                    },
                    {
                        "name": "__Host-next-auth.csrf-token",
                        "value": self.csrf_token,
                        "domain": "labs.google",
                        "path": "/",
                        "secure": True,
                        "httpOnly": True,
                        "sameSite": "Lax",
                    },
                    {
                        "name": "__Secure-next-auth.callback-url",
                        "value": "https%3A%2F%2Flabs.google%2Ffx%2Fpt%2Ftools%2Fflow",
                        "domain": "labs.google",
                        "path": "/",
                        "secure": True,
                        "httpOnly": True,
                        "sameSite": "Lax",
                    },
                    {
                        "name": "email",
                        "value": "viniciofrick%40gmail.com",
                        "domain": "labs.google",
                        "path": "/",
                        "secure": False,
                        "httpOnly": True,
                        "sameSite": "None",
                    },
                    {
                        "name": "EMAIL",
                        "value": "%22viniciofrick%40gmail.com%22",
                        "domain": "labs.google",
                        "path": "/",
                        "secure": False,
                        "httpOnly": False,
                        "sameSite": "None",
                    },
                ]
            )

            self._page = await self._context.new_page()

            # ─ Navigate to Flow gallery ─
            logger.info("FlowBrowserEngine: Loading gallery %s", FLOW_GALLERY)
            await self._page.goto(
                FLOW_GALLERY, wait_until="domcontentloaded", timeout=timeout_s * 1000
            )
            await asyncio.sleep(8)

            # ─ Enter a project + switch to image mode (with retry) ─
            entered = await self._enter_project()
            if not entered:
                raise RuntimeError("Could not enter any Flow project")

            # ─ Verify prompt input exists before marking ready ─
            # New Flow UI uses div[role="textbox"] instead of textarea
            prompt_input = self._page.locator('[role="textbox"]')
            if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
                # Fallback: try old textarea selector
                prompt_input = self._page.locator("#PINHOLE_TEXT_AREA_ELEMENT_ID")
                if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
                    raise RuntimeError(
                        f"Prompt input not found after setup. URL={self._page.url}"
                    )

            # ─ Install fetch interceptor ─
            await self._page.evaluate(_FETCH_INTERCEPTOR)

            self._http_client = httpx.AsyncClient(timeout=60.0, follow_redirects=True)
            self._ready = True
            logger.info("FlowBrowserEngine: Ready — project=%s", self._project_id)
            return True

        except Exception as e:
            logger.error("FlowBrowserEngine: Start failed — %s", e)
            await self.stop()
            return False

    async def _enter_project(self) -> bool:
        """Enter a Flow project using robust selectors, with retry."""
        
        for attempt in range(3):
            logger.info(
                "FlowBrowserEngine: Entering project (attempt %d)…", attempt + 1
            )

            # Method 1: find <a> links to /project/ on the main Flow page
            clicked = await self._page.evaluate("""
                () => {
                    // Find all links that go to a project
                    const links = document.querySelectorAll('a[href*="/project/"]');
                    if (links.length > 0) {
                        links[0].click();
                        return 'link';
                    }
                    return null;
                }
            """)

            if clicked:
                logger.info("FlowBrowserEngine: Clicked existing project via %s", clicked)
                await asyncio.sleep(8)
            else:
                # No existing projects - click "Novo projeto" button
                logger.info(
                    "FlowBrowserEngine: No project links found, clicking 'Novo projeto'…"
                )
                new_project_btn = await self._page.evaluate("""
                    () => {
                        const btns = document.querySelectorAll('button');
                        for (const btn of btns) {
                            if (btn.textContent.includes('Novo projeto')) {
                                btn.click();
                                return true;
                            }
                        }
                        return false;
                    }
                """)
                if new_project_btn:
                    logger.info("FlowBrowserEngine: Clicked 'Novo projeto' button")
                    await asyncio.sleep(8)
                else:
                    logger.warning("FlowBrowserEngine: 'Novo projeto' button not found")

            # Verify we entered a project
            url = self._page.url
            if "/project/" in url:
                self._project_id = (
                    url.split("/project/")[-1].split("?")[0].split("#")[0]
                )
                logger.info("FlowBrowserEngine: Entered project %s", self._project_id)

                # Wait for prompt input to be ready
                await asyncio.sleep(3)
                
                # Configure model and image count
                await self._configure_generation_settings()
                
                return True
            else:
                logger.warning(
                    "FlowBrowserEngine: Not in project (%s), retrying…", url
                )
                # Reload main page for next attempt
                if attempt < 2:
                    await self._page.goto(
                        FLOW_GALLERY, wait_until="domcontentloaded", timeout=30000
                    )
                    await asyncio.sleep(8)

        return False

    async def _configure_generation_settings(self) -> None:
        """Configure generation settings: select Nano Banana Pro model and x1 (1 image).
        
        Only configures if not already set to Nano Banana Pro to minimize clicks.
        """
        try:
            # Check if already configured to Nano Banana Pro (check main button text)
            current_model = await self._page.evaluate("""
                () => {
                    const btns = document.querySelectorAll('button');
                    for (const btn of btns) {
                        const text = btn.textContent || '';
                        if (text.includes('Banana') && btn.offsetParent !== null) {
                            return text;
                        }
                    }
                    return '';
                }
            """)
            
            # Skip if already set to Nano Banana Pro AND x1
            if 'Pro' in current_model and 'x1' in current_model:
                logger.info("FlowBrowserEngine: Already configured to Nano Banana Pro x1, skipping")
                return
            
            # Step 1: Open main settings dropdown
            main_btn = self._page.locator('button:has-text("Nano Banana")').first
            if await main_btn.is_visible():
                await main_btn.click()
                await asyncio.sleep(0.8)
                logger.info("FlowBrowserEngine: Opened settings dropdown")
            
            # Step 2: Click x1 button
            x1_btn = self._page.locator('button:has-text("x1")').first
            if await x1_btn.is_visible():
                await x1_btn.click()
                await asyncio.sleep(0.3)
                logger.info("FlowBrowserEngine: Selected x1 (1 image)")
            
            # Step 3: Click the model dropdown (button with arrow_drop_down)
            model_dropdown = self._page.locator('button:has-text("arrow_drop_down")').first
            if await model_dropdown.is_visible():
                await model_dropdown.click()
                await asyncio.sleep(0.8)
                
                # Step 4: Click Nano Banana Pro option
                pro_btn = self._page.locator('button:has-text("Nano Banana Pro")').first
                if await pro_btn.is_visible():
                    await pro_btn.click()
                    await asyncio.sleep(0.5)
                    logger.info("FlowBrowserEngine: Selected Nano Banana Pro model")
            
            # Close any open dropdowns
            await self._page.keyboard.press('Escape')
            await asyncio.sleep(0.3)
                
        except Exception as e:
            logger.warning("FlowBrowserEngine: Could not configure settings: %s", e)

    async def stop(self):
        """Shutdown browser."""
        self._ready = False
        try:
            if self._http_client:
                await self._http_client.aclose()
        except Exception:
            pass
        self._http_client = None
        try:
            if self._browser:
                await self._browser.close()
        except Exception:
            pass
        try:
            if self._pw:
                await self._pw.stop()
        except Exception:
            pass
        self._browser = self._context = self._page = None
        self._pw = None
        logger.info("FlowBrowserEngine: Stopped")

    # ── Session recovery ──────────────────────────────────────────────────────

    async def _refresh_session(self):
        """Refresh reCAPTCHA context by navigating back to gallery and re-entering project."""
        logger.info("FlowBrowserEngine: Refreshing session (page reload)…")
        await self._page.goto(
            FLOW_GALLERY, wait_until="domcontentloaded", timeout=30000
        )
        await asyncio.sleep(8)
        entered = await self._enter_project()
        if not entered:
            raise RuntimeError("Failed to re-enter project after session refresh")
        # New Flow UI uses div[role="textbox"] instead of textarea
        prompt_input = self._page.locator('[role="textbox"]')
        if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
            prompt_input = self._page.locator("#PINHOLE_TEXT_AREA_ELEMENT_ID")
            if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
                raise RuntimeError("Prompt input not found after session refresh")
        await self._page.evaluate(_FETCH_INTERCEPTOR)
        logger.info("FlowBrowserEngine: Session refreshed ✓")

    async def _full_restart(self):
        """Full browser restart for a completely clean reCAPTCHA state."""
        logger.info("FlowBrowserEngine: Full browser restart…")
        self._ready = False
        await self.stop()
        ok = await self.start()
        if not ok:
            raise RuntimeError("Failed to restart FlowBrowserEngine")
        logger.info("FlowBrowserEngine: Full restart complete ✓")

    async def _human_like_behavior(self):
        """Simulate human-like mouse/scroll behavior to maintain reCAPTCHA score."""
        try:
            for _ in range(random.randint(2, 4)):
                x = random.randint(200, 1000)
                y = random.randint(100, 700)
                await self._page.mouse.move(x, y, steps=random.randint(5, 15))
                await asyncio.sleep(random.uniform(0.1, 0.3))
            scroll_y = random.randint(-150, 150)
            await self._page.mouse.wheel(0, scroll_y)
            await asyncio.sleep(random.uniform(0.2, 0.5))
        except Exception:
            pass  # non-critical

    @staticmethod
    def _is_recaptcha_error(error_msg: str) -> bool:
        """Check if an error is reCAPTCHA-related (needs session recovery)."""
        indicators = ["recaptcha", "PERMISSION_DENIED", "captcha"]
        return any(ind.lower() in error_msg.lower() for ind in indicators)

    # ── Generation ─────────────────────────────────────────────────────────────

    async def set_reference_image(self, media_id: str, weight: float = 0.85) -> None:
        """Set a reference image (by mediaGenerationId) for the next generation.

        For Coverage Shots: generate the master angle first, then call this with
        its media_id before generating subsequent angles to maintain scene fidelity.

        Args:
            media_id: The mediaGenerationId from a previous generation response.
            weight: Influence weight of the reference image (0.0–1.0). Default 0.85.
        """
        await self._page.evaluate(
            "(args) => { window.__flowReferenceMediaId = args.mediaId; window.__flowReferenceWeight = args.weight; }",
            {"mediaId": media_id, "weight": weight},
        )
        logger.info(
            "FlowBrowserEngine: Reference image set — media_id=%s… weight=%.2f",
            media_id[:20],
            weight,
        )

    async def clear_reference_image(self) -> None:
        """Clear the reference image so the next generation is text-only."""
        await self._page.evaluate("() => { window.__flowReferenceMediaId = null; }")
        logger.info("FlowBrowserEngine: Reference image cleared")

    GENERATION_COOLDOWN = int(os.getenv("FLOW_GENERATION_COOLDOWN", "4"))
    BATCH_COOLDOWN = int(os.getenv("FLOW_BATCH_COOLDOWN", "2"))  # Shorter cooldown within batch
    MAX_RETRIES = 4
    RETRY_BACKOFF = [5, 10, 20, 30]  # seconds to wait before each retry
    DEFAULT_BATCH_SIZE = 10  # Max images per batch (Flow supports up to 12)

    async def generate(
        self,
        prompt: str,
        output_path: Path | str,
        aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT",
        seed: int | None = None,
    ) -> Path:
        """Generate an image using Flow (Nano Banana Pro / GEM_PIX_2)."""
        async with self._lock:
            return await self._generate_impl(
                prompt, Path(output_path), aspect_ratio, seed
            )

    async def generate_batch(
        self,
        items: list[tuple[str, Path | str]],
        aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT",
        on_progress: Optional[callable] = None,
    ) -> list[Path]:
        """Generate multiple images in a single session with minimal cooldown.
        
        Args:
            items: List of (prompt, output_path) tuples.
            aspect_ratio: Aspect ratio for all images.
            on_progress: Optional callback(current, total) for progress updates.
            
        Returns:
            List of Paths to generated images.
        """
        async with self._lock:
            results = []
            total = len(items)
            
            for idx, (prompt, output_path) in enumerate(items):
                try:
                    path = await self._generate_impl_fast(
                        prompt, Path(output_path), aspect_ratio, None
                    )
                    results.append(path)
                    
                    if on_progress:
                        on_progress(idx + 1, total)
                    
                    # Short cooldown within batch (not the full cooldown)
                    if idx < total - 1:
                        await self._human_like_behavior()
                        await asyncio.sleep(self.BATCH_COOLDOWN)
                        
                except Exception as e:
                    logger.error(
                        "FlowBrowserEngine: Batch item %d/%d failed: %s",
                        idx + 1, total, e
                    )
                    # Continue with remaining items
                    results.append(None)
            
            # Full cooldown at end of batch
            await asyncio.sleep(self.GENERATION_COOLDOWN)
            return results

    async def _generate_impl_fast(
        self,
        prompt: str,
        output_path: Path,
        aspect_ratio: str,
        seed: int | None,
    ) -> Path:
        """Fast generation without full cooldown (for batch mode)."""
        if not self._ready:
            raise RuntimeError("FlowBrowserEngine not started.")

        output_path.parent.mkdir(parents=True, exist_ok=True)
        logger.info(
            "FlowBrowserEngine: [BATCH] Generating → %s",
            output_path.name,
        )

        last_error = None
        for attempt in range(self.MAX_RETRIES + 1):
            if attempt > 0:
                backoff = self.RETRY_BACKOFF[
                    min(attempt - 1, len(self.RETRY_BACKOFF) - 1)
                ]
                if last_error and self._is_recaptcha_error(last_error):
                    try:
                        if attempt <= 2:
                            await self._refresh_session()
                        else:
                            await self._full_restart()
                    except Exception as recovery_err:
                        raise RuntimeError(
                            f"Session recovery failed: {recovery_err}"
                        ) from recovery_err
                await asyncio.sleep(backoff)

            response = await self._submit_and_wait(prompt, aspect_ratio)

            if "error" in response:
                last_error = self._extract_error(response)
                logger.error("FlowBrowserEngine: API error — %s", last_error)
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(
                    f"Flow API error after {self.MAX_RETRIES} retries: {last_error}"
                )

            media_list = response.get("media", [])
            if not media_list:
                last_error = f"No media in response: {list(response.keys())}"
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(last_error)

            first_media = media_list[0]
            image_data = first_media.get("image", {}).get("generatedImage", {})
            fife_url = image_data.get("fifeUrl", "")

            if not fife_url:
                last_error = f"No fifeUrl in image: {list(image_data.keys())}"
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(last_error)

            dl_resp = await self._http_client.get(fife_url)
            if dl_resp.status_code != 200:
                raise RuntimeError(f"Download failed: HTTP {dl_resp.status_code}")
            output_path.write_bytes(dl_resp.content)

            media_id = image_data.get("mediaGenerationId")
            if media_id:
                self._last_media_id = media_id

            size = output_path.stat().st_size
            logger.info(
                "FlowBrowserEngine: [BATCH] Image saved → %s (%d bytes)",
                output_path.name, size
            )
            return output_path

        raise RuntimeError("Unreachable")

    async def _submit_and_wait(self, prompt: str, aspect_ratio: str) -> dict:
        """Fill prompt, submit, and wait for the API response. Returns the response dict."""
        # Re-install interceptor (in case page navigated)
        await self._page.evaluate(_FETCH_INTERCEPTOR)
        # Clear previous generation result + set aspect ratio override
        await self._page.evaluate(
            "(ar) => { window.__flowLastGeneration = null; window.__flowAspectRatio = ar; }",
            aspect_ratio,
        )

        # ─ Type prompt with slight human-like delays ─
        # New Flow UI uses div[role="textbox"] instead of textarea
        prompt_input = self._page.locator('[role="textbox"]')
        if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
            # Fallback to old textarea selector
            prompt_input = self._page.locator("#PINHOLE_TEXT_AREA_ELEMENT_ID")
            if not (await prompt_input.count() > 0 and await prompt_input.first.is_visible(timeout=5000)):
                raise RuntimeError("Prompt input not found on Flow page")
        
        await prompt_input.first.click()
        await asyncio.sleep(random.uniform(0.2, 0.5))
        
        # For contenteditable div, use keyboard to clear and type
        await self._page.keyboard.press('Control+a')
        await asyncio.sleep(0.1)
        
        # Log the first 200 chars of the prompt being sent
        logger.info("FlowBrowserEngine: Sending prompt (%d chars): %s...", 
                    len(prompt), prompt[:200].replace('\n', ' '))
        
        await self._page.keyboard.type(prompt, delay=10)
        await asyncio.sleep(random.uniform(0.8, 1.5))
        
        # Verify what was actually typed into the field
        typed_content = await self._page.evaluate("""
            () => {
                const el = document.querySelector('[role="textbox"]') || 
                           document.querySelector('#PINHOLE_TEXT_AREA_ELEMENT_ID');
                return el ? (el.innerText || el.value || '') : '';
            }
        """)
        logger.info("FlowBrowserEngine: Actually typed (%d chars): %s...", 
                    len(typed_content), typed_content[:200].replace('\n', ' '))

        # ─ Click the yellow arrow submit button ─
        clicked = await self._page.evaluate("""
            () => {
                const btns = document.querySelectorAll('button');
                for (const btn of btns) {
                    if (!btn.offsetParent) continue;
                    const text = btn.textContent.trim();
                    const rect = btn.getBoundingClientRect();
                    if (text.includes('arrow_forward') && rect.y > 550 && rect.width < 80) {
                        btn.click();
                        return true;
                    }
                }
                return false;
            }
        """)
        if not clicked:
            await ta.first.press("Enter")
        logger.info("FlowBrowserEngine: Submit triggered (clicked=%s)", clicked)

        # ─ Wait for generation response ─
        for _ in range(90):
            await asyncio.sleep(1)
            gen = await self._page.evaluate("() => window.__flowLastGeneration")
            if gen and gen.get("response"):
                # Log what prompt was actually sent to the API
                prompt_sent = await self._page.evaluate("() => window.__flowLastPromptSent || 'NOT_CAPTURED'")
                logger.info("FlowBrowserEngine: Prompt sent to API (%d chars): %s...", 
                           len(prompt_sent), prompt_sent[:300].replace('\n', ' '))
                return gen["response"]

        raise RuntimeError("Timeout waiting for Flow generation response (90s)")

    @staticmethod
    def _extract_error(response: dict) -> str:
        """Extract human-readable error from API error response."""
        err = response.get("error", {})
        if isinstance(err, dict):
            code = err.get("code", "")
            msg = err.get("message", "")
            status = err.get("status", "")
            return f"[{code}] {status}: {msg}" if code else str(err)
        return str(err)

    async def _generate_impl(
        self,
        prompt: str,
        output_path: Path,
        aspect_ratio: str,
        seed: int | None,
    ) -> Path:
        if not self._ready:
            raise RuntimeError("FlowBrowserEngine not started.")

        output_path.parent.mkdir(parents=True, exist_ok=True)
        logger.info(
            "FlowBrowserEngine: Generating → %s (aspect=%s)",
            output_path.name,
            aspect_ratio,
        )

        last_error = None
        for attempt in range(self.MAX_RETRIES + 1):
            if attempt > 0:
                backoff = self.RETRY_BACKOFF[
                    min(attempt - 1, len(self.RETRY_BACKOFF) - 1)
                ]

                # ── reCAPTCHA recovery: refresh session or full restart ──
                if last_error and self._is_recaptcha_error(last_error):
                    try:
                        if attempt <= 2:
                            logger.warning(
                                "FlowBrowserEngine: reCAPTCHA recovery L1 — page refresh (attempt %d/%d)",
                                attempt,
                                self.MAX_RETRIES,
                            )
                            await self._refresh_session()
                        else:
                            logger.warning(
                                "FlowBrowserEngine: reCAPTCHA recovery L2 — full restart (attempt %d/%d)",
                                attempt,
                                self.MAX_RETRIES,
                            )
                            await self._full_restart()
                    except Exception as recovery_err:
                        logger.error(
                            "FlowBrowserEngine: Recovery failed — %s", recovery_err
                        )
                        raise RuntimeError(
                            f"Session recovery failed: {recovery_err}"
                        ) from recovery_err
                else:
                    logger.warning(
                        "FlowBrowserEngine: Retry %d/%d after %ds backoff (last error: %s)",
                        attempt,
                        self.MAX_RETRIES,
                        backoff,
                        last_error,
                    )

                await asyncio.sleep(backoff)

            response = await self._submit_and_wait(prompt, aspect_ratio)
            logger.info("FlowBrowserEngine: Response keys: %s", list(response.keys()))

            # ─ Check for error ─
            if "error" in response:
                last_error = self._extract_error(response)
                logger.error("FlowBrowserEngine: API error — %s", last_error)
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(
                    f"Flow API error after {self.MAX_RETRIES} retries: {last_error}"
                )

            # ─ Extract image URL ─
            media_list = response.get("media", [])
            if not media_list:
                last_error = f"No media in response: {list(response.keys())}"
                logger.error("FlowBrowserEngine: %s", last_error)
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(last_error)

            first_media = media_list[0]
            image_data = first_media.get("image", {}).get("generatedImage", {})
            fife_url = image_data.get("fifeUrl", "")

            if not fife_url:
                last_error = f"No fifeUrl in image: {list(image_data.keys())}"
                logger.error("FlowBrowserEngine: %s", last_error)
                if attempt < self.MAX_RETRIES:
                    continue
                raise RuntimeError(last_error)

            # ─ Download image ─
            dl_resp = await self._http_client.get(fife_url)
            if dl_resp.status_code != 200:
                raise RuntimeError(f"Download failed: HTTP {dl_resp.status_code}")
            output_path.write_bytes(dl_resp.content)

            # ─ Capture mediaGenerationId for coverage shot reference ─
            media_id = (
                first_media.get("image", {})
                .get("generatedImage", {})
                .get("mediaGenerationId")
            )
            if media_id:
                self._last_media_id = media_id
                logger.debug("FlowBrowserEngine: Captured media_id=%s…", media_id[:20])

            size = output_path.stat().st_size
            logger.info(
                "FlowBrowserEngine: Image saved → %s (%d bytes)", output_path.name, size
            )

            # ─ Cooldown + human-like behavior to maintain reCAPTCHA score ─
            await self._human_like_behavior()
            await asyncio.sleep(self.GENERATION_COOLDOWN)
            return output_path

        raise RuntimeError("Unreachable")


# ─── Global singleton ─────────────────────────────────────────────────────────

_engine: Optional[FlowBrowserEngine] = None
_engine_lock = asyncio.Lock()


async def get_engine(
    session_token: str | None = None, csrf_token: str | None = None
) -> FlowBrowserEngine:
    """Get or create the global FlowBrowserEngine singleton.
    
    Credentials can be passed explicitly (from DB) or fall back to env vars.
    """
    global _engine
    async with _engine_lock:
        if _engine and _engine.is_running:
            return _engine
        session = session_token or os.getenv("GOOGLE_FLOW_SESSION_TOKEN", "")
        csrf = csrf_token or os.getenv("GOOGLE_FLOW_CSRF_TOKEN", "")
        if not session or not csrf:
            raise RuntimeError("Google Flow cookies not configured.")
        _engine = FlowBrowserEngine(session, csrf)
        if not await _engine.start():
            raise RuntimeError("Failed to start FlowBrowserEngine")
        return _engine


async def stop_engine():
    """Stop the global engine."""
    global _engine
    if _engine:
        await _engine.stop()
        _engine = None


async def generate_image(
    prompt: str,
    output_path: Path | str,
    aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT",
    seed: int | None = None,
    reference_media_id: str | None = None,
    reference_weight: float = 0.85,
) -> Path:
    """High-level: generate an image using the Flow browser engine.

    Args:
        prompt: Text prompt for image generation.
        output_path: Where to save the generated image.
        aspect_ratio: IMAGE_ASPECT_RATIO_PORTRAIT / LANDSCAPE / SQUARE.
        seed: Optional fixed seed for reproducibility.
        reference_media_id: If set, uses this mediaGenerationId as imageInputs
            reference (image-to-image). Use for Coverage Shot angles after the
            master frame — maintains scene environment consistency across angles.
        reference_weight: Influence strength of reference image (0.0–1.0).
    """
    engine = await get_engine()
    if reference_media_id:
        await engine.set_reference_image(reference_media_id, reference_weight)
    try:
        return await engine.generate(prompt, output_path, aspect_ratio, seed)
    finally:
        if reference_media_id:
            await engine.clear_reference_image()


async def generate_images_batch(
    items: list[tuple[str, Path | str]],
    aspect_ratio: str = "IMAGE_ASPECT_RATIO_PORTRAIT",
    batch_size: int = 10,
    on_progress: Optional[callable] = None,
    on_batch_complete: Optional[callable] = None,
) -> list[Path]:
    """High-level: generate multiple images in batches using the Flow browser engine.
    
    Keeps the browser session open between batches for maximum efficiency.
    
    Args:
        items: List of (prompt, output_path) tuples.
        aspect_ratio: Aspect ratio for all images.
        batch_size: Number of images per batch (default 10, max 12).
        on_progress: Optional callback(current, total) for per-image progress.
        on_batch_complete: Optional callback(batch_num, total_batches) after each batch.
        
    Returns:
        List of Paths to generated images (None for failed items).
    """
    engine = await get_engine()
    
    all_results = []
    total_items = len(items)
    total_batches = (total_items + batch_size - 1) // batch_size
    
    for batch_num in range(total_batches):
        start_idx = batch_num * batch_size
        end_idx = min(start_idx + batch_size, total_items)
        batch_items = items[start_idx:end_idx]
        
        logger.info(
            "FlowBrowserEngine: Processing batch %d/%d (%d items)",
            batch_num + 1, total_batches, len(batch_items)
        )
        
        def batch_progress(current, total):
            if on_progress:
                global_current = start_idx + current
                on_progress(global_current, total_items)
        
        batch_results = await engine.generate_batch(
            batch_items,
            aspect_ratio=aspect_ratio,
            on_progress=batch_progress,
        )
        all_results.extend(batch_results)
        
        if on_batch_complete:
            on_batch_complete(batch_num + 1, total_batches)
        
        logger.info(
            "FlowBrowserEngine: Batch %d/%d complete (%d/%d successful)",
            batch_num + 1, total_batches,
            sum(1 for r in batch_results if r is not None),
            len(batch_results)
        )
    
    return all_results