#!/usr/bin/env python3
"""
Fix OAuth setup for Google Drive RAG system
"""

import os
import json
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request

# Configuration
SCOPES = ['https://www.googleapis.com/auth/drive.file']
CREDENTIALS_FILE = 'credentials.json'
TOKEN_FILE = 'token.json'

def fix_oauth_setup():
    """Fix OAuth setup with proper redirect URIs"""
    print("🔧 Fixing OAuth Setup for Google Drive RAG")
    print("=" * 50)
    
    # Check if credentials file exists
    if not os.path.exists(CREDENTIALS_FILE):
        print(f"❌ {CREDENTIALS_FILE} not found!")
        print("\n📋 Please follow these steps:")
        print("1. Go to: https://console.cloud.google.com/")
        print("2. APIs & Services → Credentials")
        print("3. Create Credentials → OAuth 2.0 Client IDs")
        print("4. Application type: Desktop application")
        print("5. Download as 'credentials.json'")
        return False
    
    # Delete old token file if it exists
    if os.path.exists(TOKEN_FILE):
        print(f"🗑️ Removing old token file: {TOKEN_FILE}")
        os.remove(TOKEN_FILE)
    
    print(f"✅ Found {CREDENTIALS_FILE}")
    
    try:
        # Load and validate credentials
        with open(CREDENTIALS_FILE, 'r') as f:
            creds_data = json.load(f)
        
        print("✅ Credentials file is valid")
        print(f"   Client ID: {creds_data.get('client_id', 'N/A')[:20]}...")
        
        # Check if it's a desktop application
        if creds_data.get('installed'):
            print("✅ Desktop application credentials detected")
        else:
            print("⚠️ Warning: This doesn't look like desktop application credentials")
            print("   Make sure you selected 'Desktop application' when creating credentials")
        
    except json.JSONDecodeError:
        print("❌ Invalid JSON in credentials file")
        return False
    except Exception as e:
        print(f"❌ Error reading credentials: {e}")
        return False
    
    # Try authentication with different ports
    ports_to_try = [8080, 8081, 8082, 0]  # 0 means let the system choose
    
    for port in ports_to_try:
        try:
            print(f"\n🔐 Trying authentication on port {port if port > 0 else 'auto'}...")
            
            # Create flow
            flow = InstalledAppFlow.from_client_secrets_file(CREDENTIALS_FILE, SCOPES)
            
            if port == 0:
                # Let the system choose the port
                creds = flow.run_local_server(port=0)
            else:
                # Try specific port
                creds = flow.run_local_server(port=port)
            
            print("✅ Authentication successful!")
            
            # Save credentials
            with open(TOKEN_FILE, 'w') as token:
                token.write(creds.to_json())
            print(f"✅ Credentials saved to {TOKEN_FILE}")
            
            # Test the credentials
            print("\n🧪 Testing Google Drive access...")
            from googleapiclient.discovery import build
            service = build('drive', 'v3', credentials=creds)
            
            results = service.files().list(pageSize=1, fields="files(id, name)").execute()
            files = results.get('files', [])
            
            print("✅ Google Drive access successful!")
            print(f"   Found {len(files)} file(s) in your Drive")
            
            return True
            
        except Exception as e:
            error_msg = str(e).lower()
            if "redirect_uri_mismatch" in error_msg:
                print(f"❌ Port {port} failed: redirect_uri_mismatch")
                if port < 8082:  # Don't show this message for the last attempt
                    print("   Trying next port...")
                continue
            else:
                print(f"❌ Port {port} failed: {e}")
                if port < 8082:
                    print("   Trying next port...")
                continue
    
    print("\n❌ All authentication attempts failed!")
    print("\n🔧 Manual Fix Required:")
    print("1. Go to: https://console.cloud.google.com/")
    print("2. APIs & Services → Credentials")
    print("3. Edit your OAuth 2.0 Client ID")
    print("4. Add these to 'Authorized redirect URIs':")
    print("   - http://localhost:8080/")
    print("   - http://localhost:8081/")
    print("   - http://localhost:8082/")
    print("   - http://127.0.0.1:8080/")
    print("   - http://127.0.0.1:8081/")
    print("   - http://127.0.0.1:8082/")
    print("5. Save and try again")
    
    return False

def main():
    """Main function"""
    print("🚀 OAuth Fix for Google Drive RAG System")
    print("=" * 50)
    
    if fix_oauth_setup():
        print("\n🎉 OAuth setup fixed successfully!")
        print("✅ You can now run: python setup_google_drive_rag.py")
    else:
        print("\n❌ OAuth setup failed")
        print("💡 Please follow the manual fix instructions above")

if __name__ == "__main__":
    main()