#!/usr/bin/env python3
"""
notebooklm_refresh.py — Refresh NotebookLM auth using existing Chrome headless CDP.

Connects to Chrome headless on port 9223, navigates to NotebookLM, extracts fresh
cookies and auth tokens from the live browser session, and saves them to the
default CLI profile so 'nlm' commands work again.

This is the PRIMARY recovery tool — run it after any auth failure.

Usage:
  python3 scripts/notebooklm_refresh.py          # default port 9223
  python3 scripts/notebooklm_refresh.py --port 9222   # custom port
  python3 scripts/notebooklm_refresh.py --profile work  # custom profile

Requirements:
  - Chrome headless must be running on the specified CDP port
  - Chrome user-data-dir should have saved login credentials
"""
import argparse
import json
import os
import sys
import time

# Add project to path
PROJECT_ROOT = os.path.join(os.path.dirname(__file__), '..')
sys.path.insert(0, os.path.join(PROJECT_ROOT, 'integrations', 'notebooklm-mcp-cli', 'src'))

from notebooklm_tools.utils.cdp import (
    find_or_create_notebooklm_page_by_cdp_url,
    get_page_cookies,
    get_current_url,
    is_logged_in,
    get_page_html,
    extract_csrf_token,
    extract_build_label,
    extract_session_id,
    _cdp_http_base,
)
from notebooklm_tools.core.auth import AuthTokens, get_auth_manager
from notebooklm_tools.utils.browser import flatten_cookies


def check_chrome_alive(cdp_http: str) -> dict | None:
    """Check if Chrome CDP is responsive, return version info or None."""
    import httpx
    try:
        r = httpx.get(f"{cdp_http}/json/version", timeout=5)
        return r.json()
    except Exception:
        return None


def main():
    parser = argparse.ArgumentParser(description="Refresh NotebookLM auth from Chrome headless CDP")
    parser.add_argument("--port", type=int, default=9223, help="Chrome DevTools port (default: 9223)")
    parser.add_argument("--profile", type=str, default="default", help="NLM profile name (default: default)")
    args = parser.parse_args()

    cdp_port = args.port
    profile_name = args.profile
    cdp_http = _cdp_http_base(cdp_port)

    print(f"🔌 Connecting to Chrome headless CDP at port {cdp_port}...")

    # 1. Check Chrome is alive
    info = check_chrome_alive(cdp_http)
    if not info:
        print(f"❌ Cannot connect to Chrome on port {cdp_port}")
        print(f"   → Start Chrome: google-chrome --headless --remote-debugging-port={cdp_port} \\")
        print(f"                    --user-data-dir=~/.notebooklm-mcp-cli/chrome-profiles/{profile_name}/ \\")
        print(f"                    --no-first-run --no-sandbox --disable-gpu &")
        sys.exit(1)
    print(f"✅ Chrome headless: {info.get('Browser', 'unknown')}")

    # 2. Open/create NotebookLM tab
    print("📂 Opening NotebookLM page...")
    page = find_or_create_notebooklm_page_by_cdp_url(cdp_http)
    if not page:
        print("❌ Failed to create NotebookLM tab")
        sys.exit(1)

    ws_url = page.get("webSocketDebuggerUrl")
    if not ws_url:
        print("❌ No WebSocket URL for the page")
        sys.exit(1)

    # 3. Wait for page to load
    time.sleep(3)
    current_url = get_current_url(ws_url)
    print(f"🌐 Current URL: {current_url[:100]}")

    # 4. Check login status
    logged_in = is_logged_in(current_url)
    if not logged_in:
        print("⏳ Not logged in yet — waiting for redirect...")
        time.sleep(5)
        current_url = get_current_url(ws_url)
        print(f"🌐 After waiting: {current_url[:100]}")
        logged_in = is_logged_in(current_url)

    if not logged_in:
        print("⚠️  NOT logged in — page is at accounts.google.com")
        print("   → The Chrome profile has no valid session.")
        print("   → To fix: tunnel Chrome to your local machine via:")
        print(f"     ssh -N -L {cdp_port}:127.0.0.1:{cdp_port} root@<VPS_IP>")
        print("   → Then open chrome://inspect → add localhost:9223 → log in")
        print("   → After login, re-run this script.")
        sys.exit(1)

    print(f"✅ Logged in! Extracting tokens...")

    # 5. Extract cookies via CDP
    raw_cookies = get_page_cookies(ws_url)
    print(f"🍪 Cookies extracted: {len(raw_cookies)} total")

    nlm_cookies = [c for c in raw_cookies if 'google' in c.get('domain', '')]
    print(f"🎯 Google-domain cookies: {len(nlm_cookies)}")

    # 6. Extract CSRF + session + build label from page HTML
    html = get_page_html(ws_url)
    csrf = extract_csrf_token(html)
    build = extract_build_label(html)
    session_id = extract_session_id(html)

    print(f"🔑 CSRF token:    {'✅ PRESENT' if csrf else '❌ MISSING'}")
    print(f"🔑 Session ID:    {'✅ PRESENT' if session_id else '❌ MISSING'}")
    print(f"🔑 Build label:   {'✅ PRESENT' if build else '❌ MISSING'}")

    # 7. Save to profile
    manager = get_auth_manager(profile_name)
    try:
        manager.save_profile(
            cookies=raw_cookies,
            csrf_token=csrf or None,
            session_id=session_id or None,
            build_label=build or None,
            force=True,
        )
        print(f"✅ Auth tokens saved to profile '{profile_name}'!")
    except Exception as e:
        print(f"❌ Failed to save profile: {e}")
        sys.exit(1)

    print()
    print("🧪 Running verification...")

    # 8. Verify auth works
    from notebooklm_tools.services.auth import check_auth
    result = check_auth(profile=profile_name, live=True, timeout=10)
    if result.valid:
        print(f"✅ Auth check: VALID")
        print(f"   Profile: {profile_name}")
    else:
        print(f"⚠️  Auth check: {result.reason}")
        print(f"   Details: {result.details}")

    # 9. Quick functional test
    from notebooklm_tools.services.notebooks import list_notebooks
    try:
        notebooks = list_notebooks()
        print(f"📓 Notebooks found: {len(notebooks)}")
        for nb in notebooks[:3]:
            print(f"   - {nb.get('title', nb.get('id', '?'))}")
    except Exception as e:
        print(f"⚠️  API test failed (may need CSRF refresh): {e}")

    print()
    print("✅ Done. 'nlm' commands should now work.")


if __name__ == "__main__":
    main()