#!/usr/bin/env python3
"""
Resume authenticated session: dismiss passkey prompt, land on NotebookLM, save state.
"""
import asyncio, json
from pathlib import Path
from playwright.async_api import async_playwright

PROFILE_DIR = Path.home() / ".notebooklm" / "browser_profile"
STORAGE_PATH = Path.home() / ".notebooklm" / "storage_state.json"

async def main():
    async with async_playwright() as p:
        ctx = await p.chromium.launch_persistent_context(
            str(PROFILE_DIR),
            channel="chrome",
            headless=False,
            slow_mo=100,
            args=["--start-maximized", "--disable-blink-features=AutomationControlled"],
            ignore_default_args=["--enable-automation"],
        )
        page = ctx.pages[0] if ctx.pages else await ctx.new_page()

        print("-> Navigating to NotebookLM...", flush=True)
        await page.goto("https://notebooklm.google.com", wait_until="domcontentloaded")
        await page.wait_for_timeout(4000)
        print(f"   URL: {page.url}", flush=True)

        # Handle passkey enrollment speedbump
        if "passkeyenrollment" in page.url or "speedbump" in page.url:
            print("-> Passkey prompt detected — dismissing...", flush=True)
            # Try clicking "Not now" / "Skip" / "No thanks"
            for text in ["Not now", "Skip", "No thanks", "Maybe later", "Dismiss"]:
                try:
                    btn = page.get_by_role("button", name=text)
                    if await btn.count() > 0:
                        await btn.click()
                        print(f"   Clicked: {text}", flush=True)
                        await page.wait_for_timeout(3000)
                        break
                except Exception:
                    continue
            else:
                # Try pressing Escape or clicking any "not now" link
                await page.keyboard.press("Escape")
                await page.wait_for_timeout(2000)
                # Force navigate to the continue URL
                await page.goto("https://notebooklm.google.com", wait_until="domcontentloaded")
                await page.wait_for_timeout(5000)

            print(f"   URL after dismiss: {page.url}", flush=True)

        # Handle any other interstitials
        if "accounts.google.com" in page.url:
            print("-> Still on Google — force-navigating to NotebookLM...", flush=True)
            await page.goto("https://notebooklm.google.com", wait_until="networkidle", timeout=30000)
            await page.wait_for_timeout(5000)
            print(f"   URL: {page.url}", flush=True)

        if "accounts.google.com" in page.url:
            print("FAIL: Authentication didn't carry through", flush=True)
            await ctx.close()
            return

        # We're on NotebookLM — save state
        print("-> On NotebookLM! Saving state...", flush=True)
        await ctx.storage_state(path=str(STORAGE_PATH))

        state = json.loads(STORAGE_PATH.read_text())
        cookies = state.get("cookies", [])
        print(f"   Total cookies: {len(cookies)}", flush=True)
        for c in cookies:
            print(f"   {c['name']} @ {c['domain']}", flush=True)

        has_sid = any(c["name"] == "SID" for c in cookies)
        has_psid = any("PSID" in c["name"] for c in cookies)
        print(f"\nSID: {has_sid} | PSID variant: {has_psid}", flush=True)

        if has_sid or has_psid:
            print("SUCCESS", flush=True)
        else:
            print("WARNING: Key cookies missing", flush=True)

        await ctx.close()

if __name__ == "__main__":
    asyncio.run(main())