"""
Flow Token Service — uses headless Playwright to automatically extract
project_id and fresh reCAPTCHA tokens from labs.google/fx.

This service:
1. Loads labs.google/fx with the user's cookies
2. Intercepts network requests to capture the project_id
3. Extracts reCAPTCHA tokens on-demand via grecaptcha.enterprise.execute()
4. Keeps the browser alive so tokens can be refreshed without reloading

Usage:
    service = FlowTokenService(session_token, csrf_token)
    await service.start()
    project_id = service.project_id
    recaptcha_token = await service.get_fresh_recaptcha_token()
    await service.stop()
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

RECAPTCHA_SITE_KEY = "6LdsFiUsAAAAAIjVDZcuLhaHiDn5nnHVXVRQGeMV"
LABS_URL = "https://labs.google/fx/tools/image-fx"


class FlowTokenService:
    """Headless browser service that extracts Flow API credentials."""

    def __init__(
        self,
        session_token: str,
        csrf_token: str,
    ):
        self.session_token = session_token
        self.csrf_token = csrf_token

        self.project_id: Optional[str] = None
        self._recaptcha_token: Optional[str] = None
        self._browser = None
        self._context = None
        self._page = None
        self._running = False

    async def start(self, timeout_ms: int = 30000) -> dict:
        """
        Launch headless browser, navigate to labs.google/fx, and extract credentials.

        Returns dict with:
            - project_id: str or None
            - recaptcha_token: str or None
            - error: str or None
        """
        from playwright.async_api import async_playwright

        result = {"project_id": None, "recaptcha_token": None, "error": None}

        try:
            self._pw = await async_playwright().start()
            self._browser = await self._pw.chromium.launch(
                headless=True,
                args=[
                    "--no-sandbox",
                    "--disable-blink-features=AutomationControlled",
                ],
            )

            # Create context with cookies
            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": 800},
            )

            # Add cookies
            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",
                },
            ])

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

            # Intercept network requests to find project_id
            captured_project_id = None

            async def on_request(request):
                nonlocal captured_project_id
                url = request.url
                # The flow API URL contains the project_id
                match = re.search(
                    r"aisandbox-pa\.googleapis\.com/v1/projects/([^/]+)/",
                    url,
                )
                if match:
                    captured_project_id = match.group(1)
                    logger.info("Captured project_id from request: %s", captured_project_id)

                # Also check POST body for project_id
                if "aisandbox-pa.googleapis.com" in url:
                    try:
                        body = request.post_data
                        if body:
                            data = json.loads(body)
                            # Look for projectId in client_context
                            ctx = data.get("clientContext") or data.get("client_context", {})
                            pid = ctx.get("projectId") or ctx.get("project_id")
                            if pid:
                                captured_project_id = pid
                                logger.info("Captured project_id from body: %s", pid)

                            # Also grab recaptcha token if present
                            recap = ctx.get("recaptchaContext") or ctx.get("recaptcha_context", {})
                            token = recap.get("token")
                            if token:
                                self._recaptcha_token = token
                                logger.info("Captured reCAPTCHA token from request (len=%d)", len(token))
                    except Exception:
                        pass

            self._page.on("request", on_request)

            # Navigate to labs.google/fx
            logger.info("Navigating to %s...", LABS_URL)
            await self._page.goto(LABS_URL, wait_until="networkidle", timeout=timeout_ms)
            logger.info("Page loaded: %s", self._page.url)

            # Wait a moment for any background requests
            await asyncio.sleep(3)

            # Try to extract project_id from the page's JavaScript context
            if not captured_project_id:
                try:
                    captured_project_id = await self._page.evaluate("""
                        () => {
                            // Deep search: recursively look for projectId in __NEXT_DATA__
                            const nextData = document.getElementById('__NEXT_DATA__');
                            if (nextData) {
                                try {
                                    const text = nextData.textContent;
                                    // Regex search for projectId pattern
                                    const match = text.match(/"projectId"\\s*:\\s*"([a-z0-9]+)"/);
                                    if (match) return match[1];
                                } catch(e) {}
                            }
                            // Search all script tags for projectId
                            const scripts = document.querySelectorAll('script');
                            for (const s of scripts) {
                                const t = s.textContent || '';
                                const match = t.match(/projectId['"\\s:=]+['"]([a-z0-9]{10,})['"]/);
                                if (match) return match[1];
                            }
                            // Check window/global
                            if (window.__PROJECT_ID__) return window.__PROJECT_ID__;
                            return null;
                        }
                    """)
                    if captured_project_id:
                        logger.info("Found project_id in page JS: %s", captured_project_id)
                except Exception as e:
                    logger.debug("Could not extract project_id from JS: %s", e)

            # Try to get a reCAPTCHA token
            if not self._recaptcha_token:
                try:
                    self._recaptcha_token = await self._extract_recaptcha_token()
                except Exception as e:
                    logger.warning("Could not extract reCAPTCHA token: %s", e)

            # If we still don't have project_id, trigger a real generation to capture it
            if not captured_project_id:
                logger.info("project_id not found passively. Triggering a generation to capture it...")
                try:
                    captured_project_id = await self._trigger_and_capture_project_id()
                except Exception as e:
                    logger.warning("Could not trigger project_id capture: %s", e)

            self.project_id = captured_project_id
            result["project_id"] = captured_project_id
            result["recaptcha_token"] = self._recaptcha_token
            self._running = True

            logger.info(
                "Auto-setup complete: project_id=%s, has_recaptcha=%s",
                captured_project_id or "NOT FOUND",
                bool(self._recaptcha_token),
            )

        except Exception as e:
            result["error"] = str(e)
            logger.error("Auto-setup failed: %s", e)
            await self.stop()

        return result

    async def _extract_recaptcha_token(self) -> Optional[str]:
        """Execute reCAPTCHA Enterprise on the page to get a fresh token."""
        if not self._page:
            return None

        # Wait for reCAPTCHA script to load
        try:
            await self._page.wait_for_function(
                "() => typeof grecaptcha !== 'undefined' && typeof grecaptcha.enterprise !== 'undefined'",
                timeout=10000,
            )
        except Exception:
            logger.warning("grecaptcha.enterprise not found on page, trying alternative...")
            # Try waiting a bit more
            await asyncio.sleep(2)
            try:
                await self._page.wait_for_function(
                    "() => typeof grecaptcha !== 'undefined'",
                    timeout=5000,
                )
            except Exception:
                logger.warning("No grecaptcha found at all")
                return None

        token = await self._page.evaluate(f"""
            async () => {{
                try {{
                    if (typeof grecaptcha !== 'undefined' && grecaptcha.enterprise) {{
                        const token = await grecaptcha.enterprise.execute('{RECAPTCHA_SITE_KEY}', {{action: 'generate'}});
                        return token;
                    }}
                    if (typeof grecaptcha !== 'undefined' && grecaptcha.execute) {{
                        const token = await grecaptcha.execute('{RECAPTCHA_SITE_KEY}', {{action: 'generate'}});
                        return token;
                    }}
                    return null;
                }} catch(e) {{
                    return null;
                }}
            }}
        """)
        if token:
            logger.info("Extracted fresh reCAPTCHA token (len=%d)", len(token))
        return token

    async def _trigger_and_capture_project_id(self) -> Optional[str]:
        """Type a simple prompt and click generate to capture project_id from the API request."""
        if not self._page:
            return None

        try:
            # Find and fill the prompt textarea
            # ImageFX uses a contenteditable div or textarea
            prompt_selectors = [
                'textarea[aria-label*="prompt" i]',
                'textarea[placeholder*="prompt" i]',
                'div[contenteditable="true"]',
                'textarea',
                'input[type="text"]',
            ]

            filled = False
            for sel in prompt_selectors:
                try:
                    el = self._page.locator(sel).first
                    if await el.is_visible(timeout=2000):
                        await el.click(timeout=2000)
                        await el.fill("a red apple on a table", timeout=3000)
                        filled = True
                        logger.info("Filled prompt input via: %s", sel)
                        break
                except Exception:
                    continue

            if not filled:
                logger.warning("Could not find/fill prompt input")
                return None

            await asyncio.sleep(1)

            # Try to find and click a generate/submit button
            button_selectors = [
                'button[aria-label*="generate" i]',
                'button[aria-label*="create" i]',
                'button:has-text("Generate")',
                'button:has-text("Create")',
                'button[type="submit"]',
                # ImageFX often uses a send icon button
                'button[aria-label*="send" i]',
            ]

            clicked = False
            for sel in button_selectors:
                try:
                    btn = self._page.locator(sel).first
                    if await btn.is_visible(timeout=2000):
                        await btn.click(timeout=3000)
                        clicked = True
                        logger.info("Clicked generate button via: %s", sel)
                        break
                except Exception:
                    continue

            if not clicked:
                # Try pressing Enter as fallback
                logger.info("No generate button found, pressing Enter...")
                await self._page.keyboard.press("Enter")

            # Wait for the API request to be captured
            logger.info("Waiting for API request with project_id...")
            for _ in range(15):
                await asyncio.sleep(1)
                if self.project_id:
                    logger.info("Captured project_id: %s", self.project_id)
                    return self.project_id

            logger.warning("project_id not captured after triggering generation")
            return None

        except Exception as e:
            logger.warning("Error triggering generation: %s", e)
            return None

    async def get_fresh_recaptcha_token(self) -> Optional[str]:
        """Get a fresh reCAPTCHA token (call this before each API request)."""
        if not self._page or not self._running:
            return self._recaptcha_token

        try:
            token = await self._extract_recaptcha_token()
            if token:
                self._recaptcha_token = token
            return self._recaptcha_token
        except Exception as e:
            logger.warning("Failed to refresh reCAPTCHA token: %s", e)
            return self._recaptcha_token

    async def stop(self):
        """Close the browser and clean up."""
        self._running = False
        try:
            if self._browser:
                await self._browser.close()
        except Exception:
            pass
        try:
            if self._pw:
                await self._pw.stop()
        except Exception:
            pass
        self._browser = None
        self._context = None
        self._page = None
        logger.info("Token service stopped")


async def auto_setup_flow_credentials(
    session_token: str,
    csrf_token: str,
) -> dict:
    """
    One-shot auto-setup: extract project_id and recaptcha_token, then close browser.

    Returns:
        {
            "project_id": str | None,
            "recaptcha_token": str | None,
            "error": str | None,
        }
    """
    service = FlowTokenService(session_token, csrf_token)
    result = await service.start(timeout_ms=45000)
    await service.stop()
    return result