google_sheets_utils.py

# google_sheets_utils.py
# Filepath: ai-email-assistant/google_sheets_utils.py
# Manages interactions with Google Sheets for profile and template storage

import gspread
import json
import os
import logging
from google.oauth2.service_account import Credentials
import uuid
import time
from google.auth.exceptions import TransportError
from dotenv import load_dotenv

logging.basicConfig(level=logging.INFO)

# Constants
SCOPES = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive']
SPREADSHEET_NAME = "AI Email Assistant Data"
WORKSHEET_NAMES = ["Sender Profiles", "Receiver Profiles", "Email Templates"]
HEADERS = {
    "Sender Profiles": ["Type", "ID", "Name", "Email", "Position", "Company", "Context"],
    "Receiver Profiles": ["Type", "ID", "Name", "Email", "Position", "Company", "Context"],
    "Email Templates": ["Type", "ID", "Name", "Subject", "Body"]
}
SAMPLE_DATA = {
    "Sender Profiles": ["Sender Profile", "PROF_1", "John Doe", "john.doe@example.com", "Manager", "Example Corp", "Sample context for John Doe"],
    "Receiver Profiles": ["Receiver Profile", "PROF_2", "Jane Smith", "jane.smith@example.com", "Developer", "Example Inc", "Sample context for Jane Smith"],
    "Email Templates": ["Template", "TEMP_1", "Welcome Email", "Welcome to our service", "Hello, welcome to our service. We are glad to have you."]
}

def load_config():
    load_dotenv()
    return {
        'google_sheets_credentials_file': os.getenv('GOOGLE_SHEETS_CREDENTIALS_FILE')
    }

def get_credentials():
    config = load_config()
    if not config:
        return None
    
    creds_json = config.get('google_sheets_credentials_file')
    if not creds_json:
        logging.error("Google Sheets credentials file path not found in .env file")
        return None
    
    try:
        creds_dict = json.loads(creds_json)
        return Credentials.from_service_account_info(creds_dict, scopes=SCOPES)
    except Exception as e:
        logging.error("Error creating credentials: %s", str(e))
        return None

def get_client():
    creds = get_credentials()
    if not creds:
        logging.error("Failed to obtain credentials")
        return None
    return gspread.authorize(creds)

def get_or_create_spreadsheet():
    gc = get_client()
    if not gc:
        logging.error("Failed to get Google Sheets client")
        return None
    
    max_retries = 3
    retry_delay = 5  # seconds

    for attempt in range(max_retries):
        try:
            sh = gc.open(SPREADSHEET_NAME)
            logging.info(f"Opened existing spreadsheet: {SPREADSHEET_NAME}")
            return sh
        except gspread.SpreadsheetNotFound:
            try:
                sh = gc.create(SPREADSHEET_NAME)
                logging.info(f"Created new spreadsheet: {SPREADSHEET_NAME}")
                return sh
            except Exception as e:
                logging.error(f"Failed to create spreadsheet: {str(e)}")
                return None
        except (gspread.exceptions.APIError, TransportError) as e:
            logging.error(f"API Error (attempt {attempt + 1}/{max_retries}): {str(e)}")
            if attempt < max_retries - 1:
                logging.info(f"Retrying in {retry_delay} seconds...")
                time.sleep(retry_delay)
            else:
                logging.error("Max retries reached. Unable to connect to Google Sheets.")
                return None
    
    # Ensure all required worksheets exist
    existing_worksheets = [worksheet.title for worksheet in sh.worksheets()]
    for worksheet_name in WORKSHEET_NAMES:
        if worksheet_name not in existing_worksheets:
            try:
                worksheet = sh.add_worksheet(title=worksheet_name, rows=1, cols=len(HEADERS[worksheet_name]))
                worksheet.append_row(HEADERS[worksheet_name])
                worksheet.append_row(SAMPLE_DATA[worksheet_name])
                logging.info(f"Created worksheet: {worksheet_name}")
            except Exception as e:
                logging.error(f"Failed to create worksheet {worksheet_name}: {str(e)}")
    
    return sh

def get_profiles_from_sheet(worksheet):
    try:
        records = worksheet.get_all_records()
        return records
    except Exception as e:
        logging.error("Error getting profiles from worksheet: %s", str(e))
        return []

def save_profile(worksheet, profile_data):
    try:
        worksheet.append_row(profile_data)
        logging.info("Profile saved successfully")
        return True
    except Exception as e:
        logging.error("Error saving profile: %s", str(e))
        return False

def delete_profile_from_sheet(worksheet, profile_id):
    try:
        cell = worksheet.find(profile_id)
        worksheet.delete_rows(cell.row)  # Use delete_rows instead of delete_row
        logging.info(f"Profile {profile_id} deleted successfully")
        return True
    except Exception as e:
        logging.error(f"Error deleting profile {profile_id}: {str(e)}")
        return False

def get_templates(worksheet):
    try:
        records = worksheet.get_all_records()
        return records
    except Exception as e:
        logging.error(f"Error getting templates: {str(e)}")
        return []

def save_template(worksheet, template_id, name, template_type, subject, body):
    try:
        template_data = [template_type, template_id, name, subject, body]
        worksheet.append_row(template_data)
        logging.info(f"Template saved successfully with ID: {template_id}")
        return f"Template saved successfully with ID: {template_id}"
    except Exception as e:
        logging.error(f"Error saving template: {str(e)}")
        return f"Error saving template: {str(e)}"

def delete_template_from_sheet(worksheet, template_id):
    try:
        cell = worksheet.find(template_id)
        worksheet.delete_rows(cell.row)  # Use delete_rows instead of delete_row
        logging.info(f"Template {template_id} deleted successfully")
        return f"Template {template_id} deleted successfully"
    except Exception as e:
        logging.error(f"Error deleting template {template_id}: {str(e)}")
        return f"Error deleting template {template_id}: {str(e)}"

def update_profile_in_sheet(worksheet, profile_id, updated_data):
    try:
        cell = worksheet.find(profile_id)
        for idx, value in enumerate(updated_data):
            worksheet.update_cell(cell.row, idx + 1, value)
        logging.info(f"Profile {profile_id} updated successfully")
        return True
    except Exception as e:
        logging.error(f"Error updating profile {profile_id}: {str(e)}")
        return False

def get_templates_from_sheet(worksheet):
    try:
        # Get all records from the worksheet
        records = worksheet.get_all_records()
        
        # Convert records to a list of dictionaries
        templates = []
        for record in records:
            template = {
                "Name": record.get("Name", ""),
                "Type": record.get("Type", ""),
                "ID": record.get("ID", ""),
                "Subject": record.get("Subject", ""),
                "Body": record.get("Body", "")
            }
            templates.append(template)
        
        return templates
    
    except Exception as e:
        print(f"An error occurred while fetching templates: {str(e)}")
        return []

def get_profile_summaries(worksheet, profile_type):
    try:
        profiles = worksheet.get_all_records()
        return [f"{profile['Name']} ({profile_type})" for profile in profiles]
    except Exception as e:
        logging.error(f"Error getting profile summaries: {str(e)}")
        return []

def get_template_summaries(worksheet):
    try:
        records = worksheet.get_all_records()
        summaries = [{"Name": record["Name"], "ID": record["ID"]} for record in records]
        return summaries
    except Exception as e:
        logging.error(f"Error getting template summaries: {str(e)}")
        return []

# Additional utility functions can be added here as needed

def generate_template_id():
    return f"TEMP_{uuid.uuid4().hex[:8].upper()}"

def generate_profile_id():
    return f"PROF_{uuid.uuid4().hex[:8].upper()}"

def update_template(worksheet, template_id, name, template_type, subject, body):
    try:
        cell = worksheet.find(template_id)
        row = cell.row
        worksheet.update(f'A{row}:E{row}', [[template_type, template_id, name, subject, body]])
        logging.info(f"Template {template_id} updated successfully")
        return True
    except Exception as e:
        logging.error(f"Error updating template {template_id}: {str(e)}")
        return False

def get_profile_by_id(worksheet, profile_id):
    try:
        cell = worksheet.find(profile_id)
        row = worksheet.row_values(cell.row)
        profile_data = dict(zip(HEADERS[worksheet.title], row))
        return profile_data
    except Exception as e:
        logging.error(f"Error fetching profile {profile_id}: {str(e)}")
        return None

def get_template_by_id(worksheet, template_id):
    try:
        cell = worksheet.find(template_id)
        row = worksheet.row_values(cell.row)
        template_data = dict(zip(HEADERS[worksheet.title], row))
        return template_data
    except Exception as e:
        logging.error(f"Error fetching template {template_id}: {str(e)}")
        return None

def update_template_in_sheet(worksheet, template_id, updated_data):
    try:
        cell = worksheet.find(template_id)
        for idx, value in enumerate(updated_data):
            worksheet.update_cell(cell.row, idx + 1, value)
        logging.info(f"Template {template_id} updated successfully")
        return True
    except Exception as e:
        logging.error(f"Error updating template {template_id}: {str(e)}")
        return False

if __name__ == "__main__":
    # Test the functions
    print("Testing Google Sheets integration...")
    sh = get_or_create_spreadsheet()
    if sh:
        print(f"Successfully connected to spreadsheet: {sh.title}")
        print("Worksheets:", [worksheet.title for worksheet in sh.worksheets()])
    else:
        print("Failed to connect to Google Sheets")