#!/usr/bin/env python3
"""
Use the already-authenticated browser profile to navigate to NotebookLM
and capture the correct cookies/storage 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=50,
            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...")
        await page.goto("https://notebooklm.google.com", wait_until="networkidle", timeout=60000)
        await page.wait_for_timeout(5000)
        
        url = page.url
        print(f"  Current URL: {url}")
        await page.screenshot(path="/tmp/nlm_capture1.png")

        if "accounts.google.com" in url:
            print("❌ Still on login page — not authenticated")
            await ctx.close()
            return

        print(f"✅ On NotebookLM! Waiting for full load...")
        await page.wait_for_timeout(5000)
        await page.screenshot(path="/tmp/nlm_capture2.png")

        # Save storage state
        await ctx.storage_state(path=str(STORAGE_PATH))
        state = json.loads(STORAGE_PATH.read_text())
        cookies = state.get("cookies", [])
        
        print(f"\n=== Cookies ({len(cookies)} total) ===")
        for c in cookies:
            print(f"  {c['name']} @ {c['domain']}")
        
        notebooklm_cookies = [c for c in cookies if "google" in c["domain"]]
        has_sid = any(c["name"] == "SID" for c in cookies)
        has_hsid = any(c["name"] == "HSID" for c in cookies)
        has_gaps = any(c["name"] == "__Secure-1PSID" or c["name"] == "__Secure-3PSID" for c in cookies)
        
        print(f"\nHas SID: {has_sid}")
        print(f"Has HSID: {has_hsid}")
        print(f"Has __Secure-*PSID: {has_gaps}")
        
        if has_sid or has_hsid or has_gaps:
            print("\nSUCCESS — state saved!")
        else:
            print("\nWARNING: Key cookies missing — login may have failed")
        
        await ctx.close()

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