"""Token Factory: capture BotGuard tokens from Chrome for Pro/Thinking model switching.

Chrome generates BotGuard tokens that validate model selection headers.
The Token Factory intercepts the XHR request before it's sent, captures the
token, and lets Python replay the request with the desired model header.

All methods are sync (CDP uses sync websocket library). Called from async
client via run_in_executor.
"""

from __future__ import annotations

import atexit
import os
import shutil
import signal
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

from loguru import logger

from .cdp import (
    _write_port_map,
    find_available_port,
    find_existing_gemcli_chrome,
    find_or_create_gemini_page,
    get_chrome_path,
    get_current_url,
    get_page_cookies,
    navigate_to_url,
    run_js,
)
from .constants import (
    CHROME_HEADERS,
    CHROME_USER_AGENT,
    compute_xbv,
)
from .exceptions import AuthError, TokenFactoryError

GEMINI_APP_URL = "https://gemini.google.com/app"

# XHR interceptor JS — hooks XMLHttpRequest to capture StreamGenerate requests.
# Captures the URL, headers (including BotGuard token set by JS), and body,
# then aborts the original request so Chrome doesn't send it.
XHR_INTERCEPTOR_JS = """
(() => {
    window.__cap = null;
    const origOpen = XMLHttpRequest.prototype.open;
    const origSend = XMLHttpRequest.prototype.send;
    const origSetHeader = XMLHttpRequest.prototype.setRequestHeader;

    XMLHttpRequest.prototype.open = function(m, u, ...r) {
        this.__u = u; this.__h = {};
        return origOpen.call(this, m, u, ...r);
    };
    XMLHttpRequest.prototype.setRequestHeader = function(n, v) {
        if (this.__h) this.__h[n] = v;
        return origSetHeader.call(this, n, v);
    };
    XMLHttpRequest.prototype.send = function(body) {
        if (this.__u?.includes('StreamGenerate') && !window.__cap) {
            window.__cap = {url: this.__u, headers: this.__h, body: body};
            this.abort();
            return;
        }
        return origSend.call(this, body);
    };
})()
"""

# JS to type a prompt into the Gemini input field
TYPE_PROMPT_JS = """
(() => {{
    const input = document.querySelector('div[contenteditable="true"]');
    if (input) {{
        input.focus(); input.click();
        document.execCommand('selectAll');
        document.execCommand('delete');
        document.execCommand('insertText', false, {prompt!r});
    }}
}})()
"""

# JS to click the send button
CLICK_SEND_JS = """
(() => {
    const btn = document.querySelector('button[aria-label="Send message"]')
        || document.querySelector('button[data-test-id="send-button"]');
    if (btn) btn.click();
})()
"""

# JS to check if the input field is ready
INPUT_READY_JS = '!!document.querySelector(\'div[contenteditable="true"]\')'

# Map tool_id integers to their names in Chrome's Tools drawer UI.
# These must match the text labels in Gemini's .toolbox-drawer-item-list-button elements.
TOOL_ID_TO_NAME = {
    14: "Images",   # IMAGE_TOOL_ID
    11: "Videos",   # VIDEO_TOOL_ID
    21: "Music",    # MUSIC_TOOL_ID
}

# JS to activate a specific tool in Chrome's Tools drawer.
# Opens the drawer, finds the button with matching text, and clicks it
# if not already checked. Returns 'already_active', 'activated', or 'not_found'.
ACTIVATE_TOOL_JS = """
((toolName) => {{
    // Check if tool is already active (deselect chip visible in chat bar)
    const deselect = document.querySelector('button.toolbox-drawer-item-deselect-button');
    if (deselect && deselect.textContent.toLowerCase().includes(toolName.toLowerCase())) {{
        return 'already_active';
    }}
    // Open the Tools drawer
    const drawerBtn = document.querySelector('button.toolbox-drawer-button')
        || document.querySelector('button[aria-label="Tools"]');
    if (!drawerBtn) return 'no_drawer';
    drawerBtn.click();
    // Wait a tick for the drawer to render, then find and click the tool
    return new Promise(resolve => setTimeout(() => {{
        const items = document.querySelectorAll('button.toolbox-drawer-item-list-button');
        for (const item of items) {{
            if (item.textContent.includes(toolName)) {{
                const checked = item.getAttribute('aria-checked');
                if (checked === 'true') {{
                    // Close drawer by clicking elsewhere
                    document.body.click();
                    resolve('already_active');
                    return;
                }}
                item.click();
                resolve('activated');
                return;
            }}
        }}
        // Close drawer if tool not found
        document.body.click();
        resolve('not_found');
    }}, 300));
}})('{tool_name}')
"""

# JS to deactivate any active tool (click the × on the chip in the chat bar)
DEACTIVATE_TOOL_JS = """
(() => {
    const deselect = document.querySelector('button.toolbox-drawer-item-deselect-button');
    if (deselect) { deselect.click(); return true; }
    return false;
})()
"""


@dataclass
class CapturedRequest:
    """A captured StreamGenerate request from Chrome, ready for Python replay."""

    url: str
    body: str
    headers: dict[str, str]
    cookies: dict[str, str]
    access_token: str | None = None  # Chrome's live SNlM0e token (WIZ_global_data)


class TokenFactory:
    """Captures BotGuard tokens from Chrome for Python HTTP replay.

    Usage:
        factory = TokenFactory(chrome_profile_path="/path/to/profile", profile_name="default")
        factory.ensure_ready()
        captured = factory.capture(prompt="Hello", model=some_model, metadata=None)
        # Use captured.url, captured.body, captured.headers, captured.cookies
        # to send via httpx
        factory.shutdown()
    """

    def __init__(
        self,
        chrome_profile_path: str | None = None,
        profile_name: str = "default",
    ):
        self._chrome_profile_path = chrome_profile_path
        self._profile_name = profile_name
        self._port: int | None = None
        self._ws_url: str | None = None
        self._process = None
        self._we_launched: bool = False
        self._ready: bool = False
        self._raw_cookies: list[dict] = []
        self._tf_profile_dir: Path | None = None  # Token Factory's own isolated profile
        atexit.register(self.shutdown)
        # Handle SIGTERM/SIGINT so cleanup runs when a parent process kills us
        # (atexit does NOT fire on SIGTERM — Python uses the OS default which
        # terminates immediately).  Signal handlers can only be set from the
        # main thread; silently skip if we're in a worker thread.
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                signal.signal(sig, self._signal_handler)
            except (OSError, ValueError):
                pass

    def _signal_handler(self, signum, frame):
        """Handle SIGTERM/SIGINT by cleaning up Chrome before exit."""
        self.shutdown()
        sys.exit(128 + signum)

    def _get_tf_profile_dir(self, source_dir: Path) -> Path:
        """Get the Token Factory's isolated profile directory.

        Creates a copy of the source Chrome profile at ``<source>_tf/`` so the
        Token Factory's headless Chrome never competes for the SingletonLock
        with ``gemcli login`` or CLI commands using the source profile.

        The copy is reused across Token Factory restarts and only refreshed
        when the source profile's Cookies file is newer (i.e. the user
        re-logged in).
        """
        tf_dir = source_dir.parent / f"{source_dir.name}_tf"
        synced_sentinel = tf_dir / ".tf_synced"

        needs_sync = False
        if not tf_dir.exists():
            needs_sync = True
        else:
            # Re-sync when the source Cookies DB is newer than our last copy.
            source_cookies = source_dir / "Default" / "Cookies"
            if source_cookies.exists() and synced_sentinel.exists():
                if source_cookies.stat().st_mtime > synced_sentinel.stat().st_mtime:
                    needs_sync = True
                    logger.debug("Token Factory: source profile updated, re-syncing")
            elif source_cookies.exists() and not synced_sentinel.exists():
                needs_sync = True

        if needs_sync:
            # Kill Chrome processes still referencing the old TF profile
            if tf_dir.exists():
                self._kill_stale_chrome(tf_dir)
                shutil.rmtree(tf_dir, ignore_errors=True)

            logger.debug(f"Token Factory: syncing profile → {tf_dir}")

            # Skip Chrome runtime artifacts (sockets, lock files, dangling
            # symlinks) that exist while Chrome is running on the source profile.
            _skip = {
                "SingletonLock", "SingletonCookie",
                "SingletonSocket", "RunningChromeVersion",
            }
            shutil.copytree(
                source_dir, tf_dir,
                ignore=lambda _dir, contents: [c for c in contents if c in _skip],
                ignore_dangling_symlinks=True,
                dirs_exist_ok=True,
            )
            synced_sentinel.touch()

        return tf_dir

    def ensure_ready(self) -> None:
        """Lazy init: find/launch Chrome, navigate to Gemini, wait for input field."""
        if self._ready:
            return

        chrome_path = get_chrome_path()
        if not chrome_path:
            raise TokenFactoryError(
                "Chrome not found. Install Google Chrome for Pro/Thinking models."
            )

        if not self._chrome_profile_path:
            raise TokenFactoryError(
                "No Chrome profile configured. Run 'gemcli login' first."
            )

        # Try to find an existing gemcli Chrome instance for our profile
        # (profile-aware — won't accidentally grab Antigravity's browser).
        # If the daemon is running ('gemcli chrome start'), it's registered
        # in the port map and will be found here.  We should ALWAYS reuse
        # a running daemon — it has the authenticated session and avoids
        # launching a stale TF profile.
        existing_port, existing_url = find_existing_gemcli_chrome(
            profile_name=self._profile_name
        )

        if existing_port:
            self._port = existing_port
            logger.debug(f"Token Factory: reusing Chrome on port {existing_port}")
        else:
            # Use an isolated Token Factory profile — a copy of the login
            # profile.  This prevents SingletonLock contention with gemcli
            # login / CLI commands that may use the source profile concurrently.
            source_dir = Path(self._chrome_profile_path)
            source_dir.mkdir(parents=True, exist_ok=True)

            profile_dir = self._get_tf_profile_dir(source_dir)
            self._tf_profile_dir = profile_dir

            # Clean up any stale Chrome processes using the TF profile
            self._kill_stale_chrome(profile_dir)

            # Use a different port range (9300+) for headless TF instances
            # so they don't block `gemcli login` from using the 9222+ range
            # when searching for existing visible browsers.
            self._port = find_available_port(start=9300)
            from gemini_web_mcp_cli.core.constants import CHROME_USER_AGENT
            args = [
                chrome_path,
                f"--remote-debugging-port={self._port}",
                "--no-first-run",
                "--no-default-browser-check",
                "--disable-extensions",
                f"--user-data-dir={profile_dir}",
                "--remote-allow-origins=*",
                "--headless",
                f"--user-agent={CHROME_USER_AGENT}",
                "--disable-blink-features=AutomationControlled",
            ]
            try:
                self._process = subprocess.Popen(
                    args, stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                time.sleep(3)
            except Exception as e:
                raise TokenFactoryError(f"Failed to launch Chrome: {e}")
            self._we_launched = True
            # Register in port map so other gemcli components can discover us
            _write_port_map(self._port, self._profile_name, self._process.pid)
            logger.debug(
                f"Token Factory: launched Chrome on port {self._port} "
                f"with TF profile {profile_dir}"
            )

        # Find or create Gemini page (retry — headless Chrome may take
        # a moment for the debugger endpoint to become responsive)
        page = None
        for _ in range(5):
            page = find_or_create_gemini_page(self._port)
            if page:
                break
            time.sleep(3)
        if not page:
            raise TokenFactoryError("Failed to open Gemini page in Chrome.")

        self._ws_url = page.get("webSocketDebuggerUrl")
        if not self._ws_url:
            raise TokenFactoryError("No WebSocket URL for Gemini page.")

        # Check if logged in
        current_url = get_current_url(self._ws_url)
        if "accounts.google.com" in current_url:
            raise AuthError("Not logged in to Gemini. Run 'gemcli login' first.")

        # Navigate to /app if not already there
        if "gemini.google.com/app" not in current_url:
            navigate_to_url(self._ws_url, GEMINI_APP_URL)

        # Wait for input field
        if not self._wait_for_input():
            raise TokenFactoryError("Gemini page failed to load (input field not found).")

        # Verify the page is actually authenticated.
        # An expired session will load the Gemini UI (input field present)
        # but without SNlM0e, meaning BotGuard-signed requests will fail
        # silently with misleading errors like "can't create image".
        snlm0e = run_js(
            self._ws_url, "window.WIZ_global_data?.SNlM0e || ''"
        )
        if not snlm0e:
            logger.info(
                "Token Factory: session expired — attempting "
                "auto sign-in via CDP..."
            )
            snlm0e = self._auto_sign_in()
            if not snlm0e:
                raise AuthError(
                    "Session expired — auto sign-in failed.\n"
                    "Run 'gemcli login' to refresh credentials."
                )

        self._ready = True
        logger.debug("Token Factory: ready")

    def _auto_sign_in(self) -> str | None:
        """Attempt automatic sign-in by clicking Gemini's Sign In button via CDP.

        Chrome's profile retains the Google account at the browser level even
        when Gemini session cookies expire.  Clicking "Sign in" triggers an
        OAuth redirect that Chrome auto-completes — zero user interaction.

        If the redirect lands on a Google account-chooser page, the first
        account is selected automatically.

        Returns:
            The SNlM0e token if sign-in succeeds, or None on failure.
        """
        from .cdp import _filter_gemini_cookies

        ws = self._ws_url

        # Click "Sign in" in the Gemini UI
        clicked = run_js(ws, """
            (() => {
                const buttons = document.querySelectorAll('button');
                for (const btn of buttons) {
                    if (btn.textContent.trim() === 'Sign in') {
                        btn.click();
                        return true;
                    }
                }
                return false;
            })()
        """)
        if not clicked:
            logger.debug("Token Factory: no 'Sign in' button found")
            return None

        # Monitor the redirect chain (max ~30s)
        for attempt in range(15):
            time.sleep(2)
            try:
                url = run_js(ws, "window.location.href") or ""
            except Exception:
                continue

            # Success: back on Gemini with valid session
            if "gemini.google.com/app" in url:
                snlm0e = run_js(
                    ws, "window.WIZ_global_data?.SNlM0e || ''"
                )
                if snlm0e:
                    logger.info(
                        f"Token Factory: auto sign-in succeeded "
                        f"(SNlM0e {len(snlm0e)} chars)"
                    )
                    # Save fresh cookies to auth state on disk
                    try:
                        cookies_list = get_page_cookies(ws)
                        cookies = _filter_gemini_cookies(cookies_list)
                        if cookies.get("__Secure-1PSID"):
                            from .auth import AuthManager
                            from .profiles import ProfileManager

                            pm = ProfileManager()
                            am = AuthManager(
                                profile_manager=pm,
                                profile_name=self._profile_name,
                            )
                            am.load()
                            am.state.cookies = cookies
                            am.state.tokens.snlm0e = snlm0e
                            am.state.last_refreshed = time.time()
                            am.save()
                            logger.debug(
                                "Token Factory: saved fresh auth state to disk"
                            )
                    except Exception as e:
                        logger.debug(f"Token Factory: save auth state failed: {e}")
                    return snlm0e
                # Page loaded but no SNlM0e yet — keep waiting
                continue

            # Google account chooser — auto-select first account
            if "accounts.google.com" in url:
                try:
                    page_text = (
                        run_js(ws, "document.body.innerText.substring(0, 500)")
                        or ""
                    )
                except Exception:
                    page_text = ""

                if "password" in page_text.lower():
                    logger.warning(
                        "Token Factory: auto sign-in hit password page — "
                        "can't auto-complete"
                    )
                    return None

                # Try clicking the first account in the chooser
                run_js(ws, """
                    (() => {
                        const selectors = [
                            '[data-identifier]',
                            '[data-email]',
                            'li[data-authuser]',
                            'div[data-authuser]',
                        ];
                        for (const sel of selectors) {
                            const el = document.querySelector(sel);
                            if (el) { el.click(); return; }
                        }
                    })()
                """)
                continue

        logger.warning("Token Factory: auto sign-in timed out")
        return None

    def _activate_tool(self, tool_id: int) -> bool:
        """Activate a tool (Images/Videos/Music) in Chrome's Gemini UI.

        Opens the Tools drawer and clicks the corresponding tool toggle.
        This ensures Chrome's BotGuard-signed request body includes the
        tool activation, making generation deterministic.

        Args:
            tool_id: Tool ID (14=Images, 11=Videos, 21=Music)

        Returns:
            True if tool was activated or already active, False otherwise.
        """
        tool_name = TOOL_ID_TO_NAME.get(tool_id)
        if not tool_name:
            logger.debug(f"Token Factory: unknown tool_id {tool_id}, skipping activation")
            return False

        js = ACTIVATE_TOOL_JS.format(tool_name=tool_name)
        result = run_js(self._ws_url, js)

        if result in ("activated", "already_active"):
            logger.debug(f"Token Factory: tool '{tool_name}' {result}")
            # Give the UI a moment to update placeholder/state
            if result == "activated":
                time.sleep(0.5)
            return True

        logger.warning(
            f"Token Factory: failed to activate tool '{tool_name}' "
            f"(result: {result})"
        )
        return False

    def _deactivate_tool(self) -> None:
        """Deactivate any active tool in Chrome's Gemini UI."""
        try:
            run_js(self._ws_url, DEACTIVATE_TOOL_JS)
        except Exception:
            pass  # Non-fatal

    def capture(
        self,
        prompt: str,
        model,
        metadata=None,
        tool_id: int | None = None,
    ) -> CapturedRequest:
        """Capture a BotGuard token by triggering a request in Chrome.

        Args:
            prompt: The message to send
            model: Model object with .name and .header attributes
            metadata: ConversationMetadata for multi-turn (has .cid for conversation ID)

        Returns:
            CapturedRequest with URL, body, headers, and cookies for Python replay.
        """
        self.ensure_ready()

        # Navigate to the right conversation
        target_url = GEMINI_APP_URL
        if metadata and metadata.cid:
            target_url = f"{GEMINI_APP_URL}/{metadata.cid}"

        current = get_current_url(self._ws_url)
        if target_url not in current:
            navigate_to_url(self._ws_url, target_url)
            if not self._wait_for_input():
                raise TokenFactoryError("Gemini page failed to load after navigation.")

        # Activate the requested tool (Images/Videos/Music) in Chrome's UI
        # BEFORE typing the prompt. This ensures the BotGuard-signed request
        # body includes the tool activation data.
        tool_activated = False
        if tool_id and tool_id in TOOL_ID_TO_NAME:
            tool_activated = self._activate_tool(tool_id)

        # Install XHR interceptor
        run_js(self._ws_url, XHR_INTERCEPTOR_JS)

        # Type the REAL user prompt into Chrome's input field.
        # BotGuard cryptographically signs the request body, so the prompt
        # text MUST be the actual user prompt — any later modification
        # (prompt replacement, body rebuild) invalidates the signature and
        # causes 0 candidates. We inject via DOM to handle long/complex text
        # that would be problematic via keystroke simulation.
        escaped_prompt = prompt.replace("\\", "\\\\").replace("`", "\\`").replace("$", "\\$")
        run_js(self._ws_url, TYPE_PROMPT_JS.format(prompt=escaped_prompt))
        time.sleep(0.5)
        run_js(self._ws_url, CLICK_SEND_JS)

        # Poll for captured request
        captured = self._poll_capture()
        if not captured:
            # Retry once: navigate fresh, re-install hook
            logger.debug("Token Factory: first capture failed, retrying...")
            navigate_to_url(self._ws_url, target_url)
            if not self._wait_for_input():
                raise TokenFactoryError("Gemini page failed to load on retry.")
            run_js(self._ws_url, XHR_INTERCEPTOR_JS)
            run_js(self._ws_url, TYPE_PROMPT_JS.format(prompt=escaped_prompt))
            time.sleep(0.5)
            run_js(self._ws_url, CLICK_SEND_JS)
            captured = self._poll_capture()
            if not captured:
                raise TokenFactoryError(
                    "Failed to capture BotGuard token after retry. "
                    "Chrome may not be on gemini.google.com."
                )

        # Build the full URL
        url = "https://gemini.google.com" + captured["url"]

        # Get JS-set headers (model header, tracking headers, etc.)
        js_headers = captured.get("headers", {})

        # Build complete Chrome-like headers
        headers = dict(CHROME_HEADERS)
        headers["x-browser-validation"] = compute_xbv(CHROME_USER_AGENT)
        headers.update(js_headers)

        # Override model header to the requested model
        if model and model.header:
            headers.update(model.header)

        # Grab ALL cookies from the browser
        cookies = self._get_google_cookies()

        # Extract the live SNlM0e access token from Chrome's window global.
        # This is the correct 'at=' token for the current Chrome session, which
        # must be used when replaying batchexecute calls with Chrome's cookies.
        # Mismatching Python's cached token with Chrome's cookies causes silent
        # empty responses (0 candidates). Retry up to 5s since WIZ_global_data
        # is populated by Gemini's JS bootstrap after page load.
        chrome_access_token: str | None = None
        for attempt in range(10):
            try:
                chrome_access_token = run_js(
                    self._ws_url,
                    "window.WIZ_global_data && window.WIZ_global_data.SNlM0e",
                )
                if chrome_access_token:
                    logger.debug(
                        f"Token Factory: captured Chrome SNlM0e "
                        f"(attempt {attempt + 1}, "
                        f"{len(chrome_access_token)} chars)"
                    )
                    break
            except Exception:
                pass
            time.sleep(0.5)

        # Fallback: extract SNlM0e from page HTML source.
        # The token is server-rendered in the HTML even if JS globals
        # haven't populated yet (same technique used by auth.py).
        if not chrome_access_token:
            try:
                import re

                from .cdp import get_page_html
                from .constants import TOKEN_PATTERNS

                html = get_page_html(self._ws_url)
                match = re.search(TOKEN_PATTERNS["snlm0e"], html)
                if match:
                    chrome_access_token = match.group(1)
                    logger.debug(
                        f"Token Factory: extracted SNlM0e from "
                        f"page HTML ({len(chrome_access_token)} chars)"
                    )
            except Exception as e:
                logger.debug(f"Token Factory: HTML SNlM0e extraction failed: {e}")

        if not chrome_access_token:
            raise AuthError(
                "Token Factory: SNlM0e not found — session expired.\n"
                "Run 'gemcli login' to refresh credentials."
            )

        # Deactivate any tool we enabled so Chrome returns to default state
        if tool_activated:
            self._deactivate_tool()

        # Navigate back to /app to leave Chrome in a clean state for next capture
        # (the aborted XHR may leave the page in a weird state)
        navigate_to_url(self._ws_url, GEMINI_APP_URL)

        return CapturedRequest(
            url=url,
            body=captured["body"],
            headers=headers,
            cookies=cookies,
            access_token=chrome_access_token,
        )

    def shutdown(self) -> None:
        """Clean up Chrome and resources."""
        if self._process and self._we_launched:
            try:
                self._process.terminate()
                self._process.wait(timeout=5)
            except Exception:
                try:
                    self._process.kill()
                except Exception:
                    pass
            self._process = None
        # Remove SingletonLock from the TF profile (what Chrome actually used)
        lock_dir = self._tf_profile_dir or (
            Path(self._chrome_profile_path) if self._chrome_profile_path else None
        )
        if lock_dir:
            (lock_dir / "SingletonLock").unlink(missing_ok=True)
        self._ready = False
        self._ws_url = None
        logger.debug("Token Factory: shutdown")

    @staticmethod
    def _kill_stale_chrome(profile_dir: Path) -> None:
        """Kill ALL Chrome processes using the given profile directory."""
        try:
            result = subprocess.run(
                ["pgrep", "-f", f"--user-data-dir={profile_dir}"],
                capture_output=True, text=True, timeout=5,
            )
            pids = [
                int(p) for p in result.stdout.strip().split("\n") if p.strip()
            ]
            for pid in pids:
                try:
                    os.kill(pid, signal.SIGTERM)
                    logger.debug(f"Token Factory: killing Chrome PID {pid}")
                except OSError:
                    pass
            if pids:
                # Wait for graceful exit (up to 5s), then SIGKILL survivors
                for _ in range(10):
                    alive = []
                    for pid in pids:
                        try:
                            os.kill(pid, 0)  # probe — raises OSError if dead
                            alive.append(pid)
                        except OSError:
                            pass
                    if not alive:
                        break
                    time.sleep(0.5)
                else:
                    for pid in alive:
                        try:
                            os.kill(pid, signal.SIGKILL)
                            logger.debug(
                                f"Token Factory: force-killed Chrome PID {pid}"
                            )
                        except OSError:
                            pass
                    time.sleep(1)
        except Exception as e:
            logger.debug(f"Token Factory: stale Chrome check skipped: {e}")
        # Remove stale SingletonLock regardless
        lock = profile_dir / "SingletonLock"
        lock.unlink(missing_ok=True)

    def _wait_for_input(self, timeout: int = 20) -> bool:
        """Wait for Gemini's input field to appear."""
        for _ in range(timeout):
            if run_js(self._ws_url, INPUT_READY_JS):
                return True
            time.sleep(1)
        return False

    def _poll_capture(self, timeout: int = 20) -> dict | None:
        """Poll window.__cap for a captured request."""
        for _ in range(timeout):
            time.sleep(1)
            result = run_js(self._ws_url, "window.__cap")
            if result:
                return result
        return None

    def get_raw_cookies(self) -> list[dict]:
        """Get the raw CDP cookie list (with domain/path info) from the last capture."""
        return self._raw_cookies

    def _get_google_cookies(self) -> dict[str, str]:
        """Get all Google/Gemini cookies from the browser."""
        all_cookies = get_page_cookies(self._ws_url)
        # Store raw cookies with domain info for cross-domain downloads
        self._raw_cookies = all_cookies
        cookies = {}
        for c in all_cookies:
            domain = c.get("domain", "")
            if "google.com" in domain or "gemini.google.com" in domain:
                cookies[c["name"]] = c["value"]
        return cookies
