import base64
import os.path
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
import openpyxl

# If modifying these scopes, delete the file token.json.
SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']

def get_gmail_service():
    creds = None
    if os.path.exists('token.json'):
        creds = Credentials.from_authorized_user_file('token.json', SCOPES)
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file(
                'credentials.json', SCOPES)
            creds = flow.run_local_server(port=0)
        with open('token.json', 'w') as token:
            token.write(creds.to_json())
    return build('gmail', 'v1', credentials=creds)

def get_email_content(service, user_id, msg_id):
    try:
        message = service.users().messages().get(userId=user_id, id=msg_id, format='full').execute()
        payload = message['payload']
        headers = payload['headers']
        
        subject = next((header['value'] for header in headers if header['name'].lower() == 'subject'), 'No Subject')
        sender = next((header['value'] for header in headers if header['name'].lower() == 'from'), 'Unknown Sender')
        date = next((header['value'] for header in headers if header['name'].lower() == 'date'), 'Unknown Date')
        
        parts = payload.get('parts', [])
        body = ""
        
        if 'body' in payload:
            body = payload['body'].get('data', '')
        elif parts:
            part = parts[0]
            body = part['body'].get('data', '')
        
        if body:
            body = base64.urlsafe_b64decode(body).decode('utf-8')
        
        return {
            'subject': subject,
            'sender': sender,
            'date': date,
            'body': body
        }
    except Exception as error:
        print(f'An error occurred: {error}')
        return None

def main():
    service = get_gmail_service()
    user_id = 'me'
    
    # Create a new workbook and select the active sheet
    workbook = openpyxl.Workbook()
    sheet = workbook.active
    sheet.title = "Emails"
    
    # Write headers
    headers = ["Subject", "Sender", "Date", "Body"]
    for col, header in enumerate(headers, start=1):
        sheet.cell(row=1, column=col, value=header)
    
    # Get the list of messages
    results = service.users().messages().list(userId=user_id, maxResults=200).execute()
    messages = results.get('messages', [])
    
    # Process each message
    for idx, message in enumerate(messages, start=2):
        email_data = get_email_content(service, user_id, message['id'])
        if email_data:
            sheet.cell(row=idx, column=1, value=email_data['subject'])
            sheet.cell(row=idx, column=2, value=email_data['sender'])
            sheet.cell(row=idx, column=3, value=email_data['date'])
            sheet.cell(row=idx, column=4, value=email_data['body'])
    
    # Save the workbook
    workbook.save("gmail_emails.xlsx")
    print("Emails have been saved to gmail_emails.xlsx")

if __name__ == '__main__':
    main()