"""
使用 Selenium 获取 NotebookLM 认证
Selenium 通过 Chrome 用户数据目录共享已有登录状态
"""

import json
import os
import time

try:
    from selenium import webdriver
    from selenium.webdriver.chrome.options import Options
    from selenium.webdriver.chrome.service import Service
except ImportError:
    print("请先安装: pip install selenium webdriver-manager")
    exit(1)

def get_auth_with_selenium():
    storage_path = os.path.expanduser("~/.notebooklm/storage_state.json")
    os.makedirs(os.path.dirname(storage_path), exist_ok=True)
    
    # Chrome 用户数据目录
    user_data_dir = r"C:\Users\asemx\AppData\Local\Google\Chrome\User Data"
    
    print("正在启动 Chrome（使用你的已登录账号）...")
    print("请确保之前的 Chrome 已完全关闭！")
    
    options = Options()
    options.add_argument(f"--user-data-dir={user_data_dir}")
    options.add_argument("--profile-directory=Default")
    options.add_argument("--no-first-run")
    options.add_argument("--no-default-browser-check")
    
    # 使用你的 Chrome
    options.binary_location = r"C:\Users\asemx\AppData\Local\Google\Chrome\Application\chrome.exe"
    
    try:
        driver = webdriver.Chrome(options=options)
    except Exception as e:
        print(f"❌ 无法启动 Chrome: {e}")
        print("\n请安装 chromedriver:")
        print("  pip install webdriver-manager")
        return None
    
    try:
        print("正在访问 NotebookLM...")
        driver.get("https://notebooklm.google.com")
        
        print("\n等待页面加载...（10秒）")
        time.sleep(10)
        
        # 获取 cookies
        cookies = driver.get_cookies()
        
        print(f"\n获取到 {len(cookies)} 个 cookies")
        
        # 转换为 Playwright storage_state 格式
        storage_state = {
            "cookies": cookies,
            "origins": []
        }
        
        # 保存
        with open(storage_path, "w", encoding="utf-8") as f:
            json.dump(storage_state, f, indent=2)
        
        print(f"\n✅ 认证信息已保存到: {storage_path}")
        
        content = json.dumps(storage_state, indent=2)
        print("\n" + "=" * 60)
        print("复制以下内容作为 HF Space Secret (NOTEBOOKLM_AUTH_JSON):")
        print("=" * 60)
        print(content)
        
        return storage_state
        
    finally:
        driver.quit()

if __name__ == "__main__":
    get_auth_with_selenium()