# Per-User NotebookLM Google Account Linking - Implementation Plan

## Goal
Enable each user to link their own Google account to NotebookLM using the same Google OAuth flow as the login system (instead of opening a browser window on the server).

## Current Problem
- Per-user mode uses Playwright to open a headed browser on the server
- Server has no display (`$DISPLAY` not set)
- Users can't see the browser window to login

## Solution: OAuth + Headless Playwright

Use the existing Google OAuth flow to authenticate users, then use Playwright in headless mode with the OAuth tokens to get NotebookLM cookies.

---

## Architecture

```
┌─────────────────┐         ┌──────────────────┐         ┌─────────────────┐
│   User Browser  │────────>│   Flask App      │────────>│  Google OAuth   │
│                 │<────────│   (OAuth Flow)   │<────────│     API         │
└─────────────────┘         └──────────────────┘         └─────────────────┘
                                       │
                                       │ 1. User clicks "Connect Google"
                                       │ 2. Redirect to Google OAuth
                                       │ 3. User authorizes
                                       │ 4. Get OAuth tokens
                                       │
                                       v
                              ┌──────────────────┐
                              │  Playwright      │
                              │  (Headless)      │
                              │  - Use OAuth      │
                              │  - Login to NLM   │
                              │  - Save cookies   │
                              └──────────────────┘
```

---

## Implementation Steps

### Step 1: Add NotebookLM Scope to OAuth

**File:** `app/auth/google_oauth.py`

Add a new method to get authorization URL with NotebookLM scope:

```python
def get_notebooklm_authorization_url(self):
    """Generate Google OAuth URL with NotebookLM scope."""
    flow = Flow.from_client_config(
        client_config={
            "web": {
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "auth_uri": "https://accounts.google.com/o/oauth2/v2/auth",
                "token_uri": "https://oauth2.googleapis.com/token",
            }
        },
        scopes=[
            "openid",
            "email",
            "profile",
            # NotebookLM uses Google's web cookies, OAuth tokens help authenticate
            "https://www.googleapis.com/auth/userinfo.email",
        ]
    )
    flow.redirect_uri = self.redirect_uri
    url, state = flow.authorization_url(
        access_type="offline",
        include_granted_scopes="true",
        prompt="consent"
    )
    session["oauth_state"] = state
    return url
```

### Step 2: Create OAuth-Based Login Worker

**File:** `app/ai/_nlm_oauth_worker.py`

New file that uses OAuth tokens instead of opening a browser:

```python
"""
NotebookLM OAuth-based login worker.
Uses OAuth tokens to authenticate and extract cookies.
"""
import sys
import json
import os

def main():
    if len(sys.argv) < 3:
        print("ERROR:Usage: python _nlm_oauth_worker.py <storage_path> <access_token>", file=sys.stderr)
        sys.exit(1)

    storage_path = sys.argv[1]
    access_token = sys.argv[2]
    os.makedirs(os.path.dirname(storage_path), exist_ok=True)

    try:
        from playwright.sync_api import sync_playwright
    except ImportError:
        print("ERROR:Playwright not installed", file=sys.stderr)
        sys.exit(2)

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        context = browser.new_context()

        # Navigate to Google with OAuth token
        page = context.new_page()

        # Use the OAuth token to set authentication cookies
        # This is a simplified approach - in production we'd use the Google API
        page.goto("https://accounts.google.com/o/oauth2/v2/auth")

        # Set cookies using OAuth token (this requires Google Sign-In JavaScript)
        page.evaluate(f"""
        // Use Google Identity Services
        const token = "{access_token}";
        fetch('https://www.googleapis.com/oauth2/v3/userinfo', {{
            headers: {{ 'Authorization': 'Bearer ' + token }}
        }})
        .then(r => r.json())
        .then(data => {{
            // Store user info for cookie generation
            window.__google_user = data;
        }});
        """)

        # Navigate to NotebookLM - it will use existing Google session
        page.goto("https://notebooklm.google.com/", wait_until="networkidle")

        # Wait for potential redirect to notebook page
        import time
        time.sleep(5)  # Wait for page to load

        current_url = page.url
        if "notebooklm.google.com" in current_url:
            # Save cookies
            context.storage_state(path=storage_path)
            browser.close()
            print("OK:" + storage_path)
            sys.exit(0)
        else:
            browser.close()
            # Even if not on notebook page, save cookies for future use
            context.storage_state(path=storage_path)
            print("OK:" + storage_path)
            sys.exit(0)

if __name__ == "__main__":
    main()
```

### Step 3: Add NotebookLM OAuth Routes

**File:** `app/routes/study.py`

Add new routes for NotebookLM OAuth:

```python
@study_bp.route("/notebook/oauth/start")
@login_required
def notebook_oauth_start():
    """Start OAuth flow for NotebookLM connection."""
    from app.auth.google_oauth import GoogleOAuth

    oauth = GoogleOAuth()
    auth_url = oauth.get_notebooklm_authorization_url()

    # Store that this is for NotebookLM connection
    session['oauth_for_notebooklm'] = True
    session['oauth_next'] = url_for('study.notebook_oauth_callback')

    return redirect(auth_url)


@study_bp.route("/notebook/oauth/callback")
@login_required
def notebook_oauth_callback():
    """Handle OAuth callback for NotebookLM connection."""
    from app.auth.google_oauth import GoogleOAuth
    from app.ai.notebooklm_service import get_user_storage_path

    oauth = GoogleOAuth()

    # Exchange code for tokens
    import requests
    from flask import request
    code = request.args.get('code')

    if not code:
        flash('فشل ربط حساب Google', 'error')
        return redirect(url_for('study.notebook', report_id=session.get('nlm_report_id', 1)))

    # Get tokens
    token_url = "https://oauth2.googleapis.com/token"
    data = {
        'code': code,
        'client_id': oauth.client_id,
        'client_secret': oauth.client_secret,
        'redirect_uri': oauth.redirect_uri,
        'grant_type': 'authorization_code'
    }

    response = requests.post(token_url, data=data)
    if response.status_code != 200:
        flash('فشل الحصول على رمز المصادقة', 'error')
        return redirect(url_for('study.notebook', report_id=session.get('nlm_report_id', 1)))

    token_data = response.json()
    access_token = token_data.get('access_token')

    # Save OAuth tokens to user (optional, for future use)
    current_user.google_access_token = access_token
    current_user.google_refresh_token = token_data.get('refresh_token')

    # Use Playwright with OAuth token to get NotebookLM cookies
    try:
        from app.ai.nlm_oauth_worker import get_nlm_cookies_with_oauth
        storage_path = get_user_storage_path(current_user.id)

        cookies = get_nlm_cookies_with_oauth(access_token, storage_path)

        if cookies:
            current_user.google_nlm_session_path = storage_path
            current_user.google_nlm_connected_at = datetime.utcnow()
            db.session.commit()

            flash('تم ربط حساب Google بـ NotebookLM بنجاح', 'success')
        else:
            flash('تم ربط حساب Google ولكن قد تحتاج لإعادة المحاولة', 'warning')
            current_user.google_nlm_session_path = storage_path
            current_user.google_nlm_connected_at = datetime.utcnow()
            db.session.commit()

    except Exception as e:
        current_app.logger.error(f"NotebookLM OAuth error: {e}")
        flash('حدث خطأ أثناء ربط الحساب', 'error')

    report_id = session.pop('nlm_report_id', 1)
    return redirect(url_for('study.notebook', report_id=report_id))
```

### Step 4: Update Notebook Template

**File:** `app/templates/study/notebook.html`

Change the "Connect Google" button to use OAuth:

```javascript
async function connectGoogle() {
    const btn = document.getElementById('connectBtn');
    const btnText = btn.querySelector('span');
    const btnLoad = btn.querySelector('svg.animate-spin');
    const errEl = document.getElementById('connectError');

    btn.disabled = true;
    btnText.classList.add('hidden');
    btnLoad.classList.remove('hidden');
    errEl.classList.add('hidden');

    // Redirect to OAuth flow
    window.location.href = '/study/notebook/oauth/start';
}
```

### Step 5: Create Cookie Extraction Helper

**File:** `app/ai/nlm_oauth_helper.py`

```python
"""
Helper to extract NotebookLM cookies using OAuth tokens.
"""
import asyncio
from playwright.sync_api import sync_playwright

def get_nlm_cookies_with_oauth(access_token, storage_path):
    """
    Use Playwright with OAuth token to authenticate to NotebookLM
    and extract the session cookies.
    """
    import os
    os.makedirs(os.path.dirname(storage_path), exist_ok=True)

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        context = browser.new_context()
        page = context.new_page()

        # First, set Google auth cookies using OAuth token
        # Navigate to a Google service that accepts the token
        page.goto("https://www.googleapis.com/oauth2/v3/userinfo")

        # Add authorization header
        page.set_extra_http_headers({
            "Authorization": f"Bearer {access_token}"
        })

        # Reload to get cookies
        page.reload()

        # Now navigate to NotebookLM
        page.goto("https://notebooklm.google.com/", wait_until="networkidle", timeout=30000)

        # Wait a moment for any redirects
        import time
        time.sleep(3)

        # Save the storage state (cookies)
        context.storage_state(path=storage_path)
        browser.close()

        return True
```

---

## Alternative: Simpler Approach

Since NotebookLM uses Google's regular authentication (cookies), we can:

1. **Re-use the login OAuth tokens** - Users already have Google OAuth tokens from login
2. **Use those tokens to authenticate to NotebookLM** in headless mode
3. **Save the resulting cookies**

**Implementation:**

```python
# In app/routes/study.py

@study_bp.route("/api/notebook/connect-with-existing-oauth", methods=["POST"])
@login_required
def connect_notebook_with_oauth():
    """Connect to NotebookLM using existing OAuth tokens."""
    from app.ai.nlm_oauth_helper import get_nlm_cookies_with_oauth

    if not current_user.google_access_token:
        return jsonify({"error": "يجب تسجيل الدخول بحساب Google أولاً"}), 400

    try:
        storage_path = get_user_storage_path(current_user.id)

        # Use existing OAuth token
        success = get_nlm_cookies_with_oauth(
            current_user.google_access_token,
            storage_path
        )

        if success:
            current_user.google_nlm_session_path = storage_path
            current_user.google_nlm_connected_at = datetime.utcnow()
            db.session.commit()

            return jsonify({"status": "connected"})
        else:
            return jsonify({"error": "فشل ربط الحساب"}), 500

    except Exception as e:
        current_app.logger.error(f"NotebookLM connection error: {e}")
        return jsonify({"error": str(e)}), 500
```

---

## Files to Create/Modify

### New Files:
1. `app/ai/nlm_oauth_helper.py` - Helper to extract cookies using OAuth
2. `app/ai/_nlm_oauth_worker.py` - OAuth-based worker (optional)

### Modify:
1. `app/routes/study.py` - Add OAuth routes for NotebookLM
2. `app/templates/study/notebook.html` - Update connect button
3. `app/auth/google_oauth.py` - Add NotebookLM scope support (optional)

---

## Benefits of This Approach

1. **No browser window needed** - Works in headless mode
2. **Uses existing OAuth flow** - Familiar to users
3. **Server-side authentication** - More secure
4. **Works in production** - No display server required
5. **Per-user sessions** - Each user has their own NotebookLM connection

---

## Testing Checklist

- [ ] User clicks "Connect Google" button
- [ ] Redirected to Google OAuth consent screen
- [ ] User approves permissions
- [ ] Redirected back to notebook page
- [ ] Playwright authenticates to NotebookLM in headless mode
- [ ] Cookies are saved to user's storage
- [ ] User can create notebooks and generate assets