#!/usr/bin/env python3
"""
Re-auth assafdagancos@gmail.com with expanded scopes (gmail.modify + calendar).
Runs a local callback server to capture the OAuth code automatically.
"""
import json, os, sys, threading, urllib.parse, urllib.request
from http.server import HTTPServer, BaseHTTPRequestHandler

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')
TOKEN_FILE = os.path.join(GOOGLE_AUTH_DIR, 'token-personal.json')
PORT = 8888

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

auth_code = None
server_done = threading.Event()

class CallbackHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        global auth_code
        parsed = urllib.parse.urlparse(self.path)
        params = urllib.parse.parse_qs(parsed.query)
        if 'code' in params:
            auth_code = params['code'][0]
            self.send_response(200)
            self.send_header('Content-Type', 'text/html')
            self.end_headers()
            self.wfile.write(b'<html><body><h2>Auth complete! You can close this tab.</h2></body></html>')
        else:
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'<html><body><h2>No code received.</h2></body></html>')
        server_done.set()

    def log_message(self, *args):
        pass  # suppress server logs

def main():
    with open(CREDENTIALS_FILE) as f:
        cred = json.load(f).get('installed', {})

    client_id = cred['client_id']
    client_secret = cred['client_secret']
    redirect_uri = f'http://localhost:{PORT}'

    params = {
        'client_id': client_id,
        'redirect_uri': redirect_uri,
        'response_type': 'code',
        'scope': ' '.join(SCOPES),
        'access_type': 'offline',
        'prompt': 'consent',
        'login_hint': 'assafdagancos@gmail.com',
    }
    auth_url = f"https://accounts.google.com/o/oauth2/auth?{urllib.parse.urlencode(params)}"

    # Start callback server
    httpd = HTTPServer(('', PORT), CallbackHandler)
    t = threading.Thread(target=httpd.handle_request)
    t.daemon = True
    t.start()

    print(f"AUTH_URL:{auth_url}")
    print(f"Waiting for callback on port {PORT}...")
    sys.stdout.flush()

    server_done.wait(timeout=120)

    if not auth_code:
        print("ERROR: No auth code received within timeout.")
        sys.exit(1)

    # Exchange code for tokens
    token_data = urllib.parse.urlencode({
        'code': auth_code,
        'client_id': client_id,
        'client_secret': client_secret,
        'redirect_uri': redirect_uri,
        '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_data = {
        'token': tokens['access_token'],
        'access_token': tokens['access_token'],
        'refresh_token': tokens.get('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"SUCCESS: Token saved to {TOKEN_FILE}")
    print(f"Scopes: {', '.join(SCOPES)}")

if __name__ == '__main__':
    main()
