google_drive_sync.py

import os
import json
import pickle
from datetime import datetime

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google.oauth2 import service_account
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload

# Google Drive API scopes
SCOPES = ['https://www.googleapis.com/auth/drive.file']

# ============================================================================
# Hard-coded Google Drive auth configuration
# Option 1: Service Account (recommended)
# Option 2: Refresh Token
# ============================================================================

# Option 1: Service Account JSON file (env-var friendly for Spaces)
SERVICE_ACCOUNT_FILE = os.getenv("GOOGLE_SERVICE_ACCOUNT_FILE", None)  # e.g. "path/to/service-account-key.json"

# Option 2: Refresh Token (requires OAuth flow once)
REFRESH_TOKEN = os.getenv("GOOGLE_REFRESH_TOKEN", None)
CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")

# Fallback: interactive OAuth
USE_HARDCODED_AUTH = SERVICE_ACCOUNT_FILE is not None or REFRESH_TOKEN is not None

# Real OAuth 2.0 configuration - Public client credentials for this app
# These are safe to include in public code (they're meant to be public)
CLIENT_CONFIG = {
    "installed": {
        "client_id": CLIENT_ID,
        "project_id": "patient-eval-app-439620",
        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
        "token_uri": "https://oauth2.googleapis.com/token",
        "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
        "client_secret": CLIENT_SECRET,
        "redirect_uris": ["http://localhost"]
    }
}

class GoogleDriveSync:
    def __init__(self, user_id="default"):
        self.user_id = user_id
        self.token_file = f'user_tokens/{user_id}_token.pickle'
        self.service = None
        self.folder_id = None
        
        # Ensure user tokens directory exists
        os.makedirs('user_tokens', exist_ok=True)
        
    def start_oauth_flow(self):
        """Start OAuth flow for user authentication"""
        try:
            # Use Google's Flow from the client config
            flow = InstalledAppFlow.from_client_config(CLIENT_CONFIG, SCOPES)
            
            # Run the OAuth flow with port=0 to use any available port
            creds = flow.run_local_server(port=0, open_browser=True)
            
            # Save the credentials
            with open(self.token_file, 'wb') as token:
                pickle.dump(creds, token)
                
            return True, f"Google authentication successful! You can now sync your evaluation data to your personal Google Drive."
                
        except Exception as e:
            return False, f"Authentication error: {str(e)}"
    
    def authenticate(self):
        """Authenticate with hardcoded credentials, existing token, or start new flow"""
        creds = None
        
        # Option 1: Service Account (supports env JSON for Spaces)
        service_account_json = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
        if service_account_json:
            try:
                import tempfile
                import json as json_lib
                # Persist JSON string to a temporary file
                with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                    json_lib.dump(json_lib.loads(service_account_json), f)
                    temp_file = f.name
                creds = service_account.Credentials.from_service_account_file(
                    temp_file, scopes=SCOPES)
                self.service = build('drive', 'v3', credentials=creds)
                os.unlink(temp_file)  # remove temp file
                return True, "Authenticated with Service Account (from env)"
            except Exception as e:
                return False, f"Service Account authentication error: {str(e)}"
        
        if SERVICE_ACCOUNT_FILE and os.path.exists(SERVICE_ACCOUNT_FILE):
            try:
                creds = service_account.Credentials.from_service_account_file(
                    SERVICE_ACCOUNT_FILE, scopes=SCOPES)
                self.service = build('drive', 'v3', credentials=creds)
                return True, "Authenticated with Service Account"
            except Exception as e:
                return False, f"Service Account authentication error: {str(e)}"
        
        # Option 2: hard-coded refresh token
        if REFRESH_TOKEN:
            if not CLIENT_ID or not CLIENT_SECRET:
                return False, "Refresh token provided but GOOGLE_CLIENT_ID / GOOGLE_CLIENT_SECRET are missing."
            try:
                creds = Credentials(
                    token=None,
                    refresh_token=REFRESH_TOKEN,
                    token_uri="https://oauth2.googleapis.com/token",
                    client_id=CLIENT_ID,
                    client_secret=CLIENT_SECRET
                )
                # Refresh to obtain a fresh access token
                creds.refresh(Request())
                self.service = build('drive', 'v3', credentials=creds)
                return True, "Authenticated with Refresh Token"
            except Exception as e:
                return False, f"Refresh Token authentication error: {str(e)}"
        
        # Option 3: previously saved token file
        if os.path.exists(self.token_file):
            try:
                with open(self.token_file, 'rb') as token:
                    creds = pickle.load(token)
                
                # Check if credentials are valid and refresh if needed
                if not creds or not creds.valid:
                    if creds and creds.expired and creds.refresh_token:
                        creds.refresh(Request())
                        # Save refreshed credentials
                        with open(self.token_file, 'wb') as token:
                            pickle.dump(creds, token)
                    else:
                        return False, "Credentials expired. Please login again."
                
                # Build the service
                self.service = build('drive', 'v3', credentials=creds)
                return True, f"Authenticated with your Google account"
                
            except Exception as e:
                # Token file corrupted, remove it
                if os.path.exists(self.token_file):
                    os.remove(self.token_file)
                return False, f"Authentication error: {str(e)}. Please login again."
                
        return False, "Please authenticate with Google first by clicking 'Login to Google Drive'"
    
    def create_or_get_folder(self, folder_name="Patient_Evaluations"):
        """Create or get the folder for storing evaluation data"""
        if not self.service:
            return None, "Not authenticated with Google Drive"
        
        try:
            # Search for existing folder
            results = self.service.files().list(
                q=f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false",
                spaces='drive',
                fields="files(id, name)"
            ).execute()
            
            items = results.get('files', [])
            if items:
                self.folder_id = items[0]['id']
                return self.folder_id, f"Found existing folder: {folder_name}"
            
            # Create new folder
            folder_metadata = {
                'name': folder_name,
                'mimeType': 'application/vnd.google-apps.folder'
            }
            
            folder = self.service.files().create(body=folder_metadata, fields='id').execute()
            self.folder_id = folder.get('id')
            return self.folder_id, f"Created new folder: {folder_name}"
            
        except Exception as e:
            return None, f"Error creating/finding folder: {str(e)}"
    
    def upload_file(self, local_file_path, drive_file_name=None):
        """Upload a file to user's Google Drive"""
        if not self.service or not self.folder_id:
            return False, "Google Drive not properly initialized"
        
        if not os.path.exists(local_file_path):
            return False, f"Local file not found: {local_file_path}"
        
        try:
            if not drive_file_name:
                drive_file_name = os.path.basename(local_file_path)
            
            # Add timestamp to filename
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            name, ext = os.path.splitext(drive_file_name)
            drive_file_name = f"{name}_{timestamp}{ext}"
            
            file_metadata = {
                'name': drive_file_name,
                'parents': [self.folder_id]
            }
            
            media = MediaFileUpload(local_file_path, resumable=True)
            
            file = self.service.files().create(
                body=file_metadata,
                media_body=media,
                fields='id,name,webViewLink'
            ).execute()
            
            # Make file shareable (optional)
            try:
                self.service.permissions().create(
                    fileId=file.get('id'),
                    body={'role': 'reader', 'type': 'anyone'}
                ).execute()
            except:
                pass  # Permission setting failed, but upload succeeded
            
            return True, {
                'file_id': file.get('id'),
                'name': file.get('name'),
                'link': file.get('webViewLink'),
                'message': f"File uploaded successfully: {drive_file_name}"
            }
            
        except Exception as e:
            return False, f"Error uploading file: {str(e)}"
    
    def upload_evaluation_data(self, csv_file_path, json_files_dir="patient_evaluations"):
        """Upload all evaluation data to user's Google Drive"""
        results = []
        
        # Upload CSV file
        if os.path.exists(csv_file_path):
            success, result = self.upload_file(csv_file_path, "patient_evaluations_master.csv")
            results.append(('CSV Master File', success, result))
        
        # Upload JSON files
        if os.path.exists(json_files_dir):
            json_files = [f for f in os.listdir(json_files_dir) if f.endswith('.json')]
            for filename in json_files[:5]:  # Limit to prevent too many uploads at once
                file_path = os.path.join(json_files_dir, filename)
                success, result = self.upload_file(file_path)
                results.append(('Evaluation File', success, result))
        
        return results

# Global instances per user
user_syncs = {}

def get_user_sync(user_id="default"):
    """Get or create GoogleDriveSync instance for user"""
    if user_id not in user_syncs:
        user_syncs[user_id] = GoogleDriveSync(user_id)
    return user_syncs[user_id]

def start_google_login(user_id="default"):
    """Start Google OAuth flow for user"""
    sync = get_user_sync(user_id)
    return sync.start_oauth_flow()

def sync_to_google_drive(user_id="default"):
    """Main function to sync evaluation data to user's Google Drive"""
    sync = get_user_sync(user_id)
    
    # Check authentication
    success, message = sync.authenticate()
    if not success:
        return False, message
    
    # Create or get folder
    folder_id, folder_message = sync.create_or_get_folder()
    if not folder_id:
        return False, folder_message
    
    # Upload evaluation data
    csv_path = "patient_evaluations/patient_evaluations_master.csv"
    results = sync.upload_evaluation_data(csv_path)
    
    success_count = sum(1 for _, success, _ in results if success)
    total_count = len(results)
    
    if success_count > 0:
        links = []
        for file_type, success, result in results:
            if success and isinstance(result, dict) and 'link' in result:
                links.append(f"• {file_type}: {result['link']}")
        
        link_text = "\n" + "\n".join(links) if links else ""
        return True, f"✅ Synced {success_count}/{total_count} files to your Google Drive!{link_text}"
    else:
        return False, "❌ Failed to sync files. Please check your authentication and try again."

def get_drive_setup_instructions():
    """Return simplified instructions for users"""
    return """
## 🔗 Save to Your Google Drive

Click "Login to Google Drive" below to connect your personal Google account.

### What happens:
1. **Secure Login**: A browser window will open for you to login to your Google account
2. **Permission Request**: You'll be asked to allow this app to save files to your Drive  
3. **Automatic Backup**: Your evaluation data will be saved to a "Patient_Evaluations" folder in your Drive
4. **Your Data**: Only you can access your files - they're saved to YOUR personal Google Drive

### First Time Setup:
- Click "Login to Google Drive" 
- Complete the Google authentication in your browser
- Return here and click "Sync to Google Drive" to backup your data

### Security:
- This app only saves evaluation files you create
- No access to your other Google Drive files
- You can revoke access anytime in your Google Account settings
"""

def check_user_authentication(user_id="default"):
    """Check if user is already authenticated"""
    sync = get_user_sync(user_id)
    success, message = sync.authenticate()
    return success, message

def get_refresh_token():
    """
    Helper: run once to fetch a refresh token for the hard-coded config.
    Complete the OAuth flow, copy the printed token into REFRESH_TOKEN.
    """
    try:
        flow = InstalledAppFlow.from_client_config(CLIENT_CONFIG, SCOPES)
        creds = flow.run_local_server(port=0, open_browser=True)
        
        if creds and creds.refresh_token:
            print("\n" + "="*60)
            print("✅ Refresh token retrieved!")
            print("="*60)
            print(f"\nPaste the following refresh_token into REFRESH_TOKEN:\n")
            print(f"REFRESH_TOKEN = \"{creds.refresh_token}\"")
            print("\n" + "="*60 + "\n")
            return creds.refresh_token
        else:
            print("❌ Could not obtain refresh_token. Please retry.")
            return None
    except Exception as e:
        print(f"❌ Error retrieving refresh_token: {str(e)}")
        return None

def auto_upload_to_drive(file_path, user_id="default", folder_name="Chatbot_Data"):
    """
    Automatically upload a file to Google Drive if user is authenticated.
    Returns (success, message) tuple. If not authenticated, returns (False, None) silently.
    This function fails silently to not interrupt normal operation.
    """
    try:
        sync = get_user_sync(user_id)
        
        # Check authentication (uses hard-coded creds if configured)
        success, message = sync.authenticate()
        if not success:
            # Log once when hard-coded auth was expected but failed
            if USE_HARDCODED_AUTH:
                print(f"⚠️ Google Drive authentication failed: {message}")
            return False, None  # Not authenticated, fail silently
        
        # Create or get folder (use custom folder name for chatbot data)
        folder_id, folder_message = sync.create_or_get_folder(folder_name)
        if not folder_id:
            return False, None  # Fail silently if folder creation fails
        
        # Upload file
        success, result = sync.upload_file(file_path)
        if success:
            return True, f"✅ Auto-uploaded to Google Drive: {result.get('name', 'file')}"
        else:
            return False, None  # Fail silently on upload error
    except Exception as e:
        # Fail silently if there's an error (e.g., network issue, import error)
        return False, None