#!/usr/bin/env python3
"""OAuth setup for additional Google accounts. Run interactively."""
import json, os, sys

WORKSPACE = os.environ.get('WORKSPACE', '/root/.openclaw/workspace')
GOOGLE_AUTH_DIR = os.path.join(WORKSPACE, 'google-auth')
CREDENTIALS_FILE = os.path.join(GOOGLE_AUTH_DIR, 'credentials.json')

SCOPES = [
    'https://www.googleapis.com/auth/gmail.readonly',
    'https://www.googleapis.com/auth/calendar.readonly',
]

def setup(account_type='personal'):
    """Generate OAuth URL and save token."""
    if account_type == 'personal':
        token_file = os.path.join(GOOGLE_AUTH_DIR, 'token-personal.json')
    else:
        token_file = os.path.join(GOOGLE_AUTH_DIR, f'token-{account_type}.json')
    
    with open(CREDENTIALS_FILE) as f:
        cred_data = json.load(f)
    
    installed = cred_data.get('installed', cred_data.get('web', {}))
    client_id = installed['client_id']
    client_secret = installed['client_secret']
    
    # Build OAuth URL
    import urllib.parse
    params = {
        'client_id': client_id,
        'redirect_uri': 'urn:ietf:wg:oauth:2.0:oob',
        'response_type': 'code',
        'scope': ' '.join(SCOPES),
        'access_type': 'offline',
        'prompt': 'consent',
    }
    auth_url = f"https://accounts.google.com/o/oauth2/auth?{urllib.parse.urlencode(params)}"
    
    print(f"\n🔐 OAuth Setup for {account_type} account")
    print(f"\n1. Open this URL in your browser:\n")
    print(auth_url)
    print(f"\n2. Sign in with the {account_type} Google account")
    print("3. Grant the requested permissions")
    print("4. Copy the authorization code and paste it below\n")
    
    code = input("Authorization code: ").strip()
    
    if not code:
        print("❌ No code provided. Aborting.")
        sys.exit(1)
    
    # Exchange code for token
    import urllib.request
    token_data = urllib.parse.urlencode({
        'code': code,
        'client_id': client_id,
        'client_secret': client_secret,
        'redirect_uri': 'urn:ietf:wg:oauth:2.0:oob',
        'grant_type': 'authorization_code',
    }).encode()
    
    req = urllib.request.Request('https://oauth2.googleapis.com/token', data=token_data)
    resp = urllib.request.urlopen(req)
    tokens = json.loads(resp.read())
    
    # Save in same format as existing token
    save_data = {
        'token': tokens['access_token'],
        'access_token': tokens['access_token'],
        'refresh_token': tokens['refresh_token'],
        'token_uri': 'https://oauth2.googleapis.com/token',
        'client_id': client_id,
        'client_secret': client_secret,
        'scopes': SCOPES,
    }
    
    with open(token_file, 'w') as f:
        json.dump(save_data, f, indent=2)
    os.chmod(token_file, 0o600)
    
    print(f"\n✅ Token saved to {token_file}")
    print(f"   Scopes: {', '.join(SCOPES)}")


if __name__ == '__main__':
    account = sys.argv[1] if len(sys.argv) > 1 else 'personal'
    if account.startswith('--account='):
        account = account.split('=')[1]
    elif account == '--account' and len(sys.argv) > 2:
        account = sys.argv[2]
    setup(account)
