app/gmail/auth.py

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
import os
import json
from loguru import logger
from config import GMAIL_TOKEN_JSON, GMAIL_CREDENTIALS_JSON, GMAIL_TOKEN_FILE, GMAIL_CREDENTIALS_FILE

SCOPES = ["https://www.googleapis.com/auth/gmail.modify"]

def authenticate():
    creds = None

    # 1. Try loading authorized user token from environment variable
    if GMAIL_TOKEN_JSON:
        try:
            logger.info("🔑 Loading Gmail token from GMAIL_TOKEN_JSON environment variable...")
            info = json.loads(GMAIL_TOKEN_JSON)
            creds = Credentials.from_authorized_user_info(info, SCOPES)
            logger.info("✅ Gmail token successfully loaded from environment variable")
        except Exception as e:
            logger.error(f"❌ Failed to load token from GMAIL_TOKEN_JSON environment variable: {e}")
            creds = None

    # 2. Try loading from local token file
    if not creds and os.path.exists(GMAIL_TOKEN_FILE):
        try:
            logger.info(f"🔑 Loading Gmail token from file: {GMAIL_TOKEN_FILE}...")
            creds = Credentials.from_authorized_user_file(GMAIL_TOKEN_FILE, SCOPES)
            logger.info("✅ Gmail token successfully loaded from file")
        except Exception as e:
            logger.error(f"❌ Failed to load token from file {GMAIL_TOKEN_FILE}: {e}")
            creds = None

    # 3. If credentials don't exist or are invalid, handle renewal/flow
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            try:
                logger.info("🔄 Gmail access token expired. Refreshing token...")
                creds.refresh(Request())
                logger.info("✅ Gmail token successfully refreshed")
                # If we're local and loaded from file, update the file
                if os.path.exists(GMAIL_TOKEN_FILE):
                    try:
                        with open(GMAIL_TOKEN_FILE, "w") as f:
                            f.write(creds.to_json())
                    except Exception as e:
                        logger.warning(f"Could not save refreshed token to file: {e}")
            except Exception as e:
                logger.error(f"❌ Failed to refresh Gmail token: {e}")
                creds = None

        if not creds:
            # Check if we are running in a headless environment (like HF Spaces)
            is_headless = os.getenv("SPACE_ID") is not None or os.getenv("PORT") is not None
            
            if is_headless:
                logger.error("❌ Gmail authorization is missing. In headless/cloud environments, please set GMAIL_TOKEN_JSON environment variable with your token.json contents!")
                raise RuntimeError("Gmail credentials are not configured. Please set GMAIL_TOKEN_JSON environment variable in Hugging Face Secrets.")

            logger.info("🌐 Authenticating via local browser flow...")
            
            # Determine how to load client secrets
            if GMAIL_CREDENTIALS_JSON:
                try:
                    secrets_info = json.loads(GMAIL_CREDENTIALS_JSON)
                    flow = InstalledAppFlow.from_client_config(secrets_info, SCOPES)
                except Exception as e:
                    logger.error(f"❌ Failed to parse client secrets from GMAIL_CREDENTIALS_JSON: {e}")
                    raise
            elif os.path.exists(GMAIL_CREDENTIALS_FILE):
                flow = InstalledAppFlow.from_client_secrets_file(GMAIL_CREDENTIALS_FILE, SCOPES)
            else:
                logger.error(f"❌ Missing client credentials. GMAIL_CREDENTIALS_FILE ({GMAIL_CREDENTIALS_FILE}) not found!")
                raise FileNotFoundError(f"Missing {GMAIL_CREDENTIALS_FILE} or GMAIL_CREDENTIALS_JSON environment variable.")

            try:
                creds = flow.run_local_server(
                    port=8080,
                    prompt="consent",
                    access_type="offline",
                    open_browser=True
                )
                
                with open(GMAIL_TOKEN_FILE, "w") as f:
                    f.write(creds.to_json())
                logger.info(f"✅ Authorization successful! token saved to {GMAIL_TOKEN_FILE}")
            except Exception as e:
                logger.error(f"❌ Authorization flow failed: {e}")
                raise

    return creds