gmail_parser.py

"""
Gmail Mail Parser
==================
Parses all Gmail emails and categorizes them by labels (Inbox, Spam, Promotions, etc.)
Extracts: Sender, Subject, Date, Snippet, Body, Labels
"""

import os
import json
import base64
from datetime import datetime
from typing import Dict, List, Optional
from email.utils import parseaddr

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
from googleapiclient.errors import HttpError

# Gmail API Scopes - readonly access to emails
SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']

# Output directory for parsed emails
OUTPUT_DIR = 'parsed_emails'


class GmailParser:
    """Gmail email parser that fetches and categorizes all emails."""
    
    def __init__(self):
        self.service = None
        self.credentials = None
        
    def authenticate(self) -> bool:
        """
        Authenticate with Gmail API using OAuth2.
        Returns True if authentication successful.
        """
        creds = None
        
        # Check for existing token
        if os.path.exists('token.json'):
            creds = Credentials.from_authorized_user_file('token.json', SCOPES)
        
        # If no valid credentials, let user log in
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                print("🔄 Refreshing expired credentials...")
                creds.refresh(Request())
            else:
                if not os.path.exists('credentials.json'):
                    print("❌ Error: 'credentials.json' not found!")
                    print("📋 Please follow the setup instructions in SETUP.md")
                    return False
                
                print("🔐 Opening browser for Gmail authentication...")
                flow = InstalledAppFlow.from_client_secrets_file('credentials.json', SCOPES)
                creds = flow.run_local_server(port=0)
            
            # Save credentials for next run
            with open('token.json', 'w') as token:
                token.write(creds.to_json())
                print("✅ Credentials saved to token.json")
        
        self.credentials = creds
        self.service = build('gmail', 'v1', credentials=creds)
        print("✅ Successfully authenticated with Gmail API!")
        return True
    
    def get_labels(self) -> Dict[str, str]:
        """Fetch all Gmail labels and return id->name mapping."""
        try:
            results = self.service.users().labels().list(userId='me').execute()
            labels = results.get('labels', [])
            return {label['id']: label['name'] for label in labels}
        except HttpError as error:
            print(f"❌ Error fetching labels: {error}")
            return {}
    
    def fetch_emails(self, max_results: int = 500, label_ids: Optional[List[str]] = None) -> List[Dict]:
        """
        Fetch emails from Gmail.
        
        Args:
            max_results: Maximum number of emails to fetch
            label_ids: Optional list of label IDs to filter by
            
        Returns:
            List of email message objects
        """
        emails = []
        page_token = None
        fetched = 0
        
        print(f"\n📧 Fetching up to {max_results} emails...")
        
        try:
            while fetched < max_results:
                # Build request parameters
                params = {
                    'userId': 'me',
                    'maxResults': min(100, max_results - fetched),
                }
                if label_ids:
                    params['labelIds'] = label_ids
                if page_token:
                    params['pageToken'] = page_token
                
                # Fetch message list
                results = self.service.users().messages().list(**params).execute()
                messages = results.get('messages', [])
                
                if not messages:
                    break
                
                # Fetch full message details for each
                for msg in messages:
                    email_data = self._get_email_details(msg['id'])
                    if email_data:
                        emails.append(email_data)
                        fetched += 1
                        
                        if fetched % 50 == 0:
                            print(f"   📥 Fetched {fetched} emails...")
                
                # Check for more pages
                page_token = results.get('nextPageToken')
                if not page_token:
                    break
                    
        except HttpError as error:
            print(f"❌ Error fetching emails: {error}")
        
        print(f"✅ Successfully fetched {len(emails)} emails!")
        return emails
    
    def _get_email_details(self, msg_id: str) -> Optional[Dict]:
        """
        Get detailed information for a single email.
        
        Args:
            msg_id: Gmail message ID
            
        Returns:
            Dictionary with email details or None if error
        """
        try:
            message = self.service.users().messages().get(
                userId='me', 
                id=msg_id, 
                format='full'
            ).execute()
            
            headers = message.get('payload', {}).get('headers', [])
            
            # Extract header values
            subject = self._get_header(headers, 'Subject') or '(No Subject)'
            from_header = self._get_header(headers, 'From') or 'Unknown'
            to_header = self._get_header(headers, 'To') or ''
            date_header = self._get_header(headers, 'Date') or ''
            
            # Parse sender name and email
            sender_name, sender_email = parseaddr(from_header)
            
            # Get email body
            body = self._get_email_body(message.get('payload', {}))
            
            # Get labels
            label_ids = message.get('labelIds', [])
            
            # Determine category
            category = self._categorize_email(label_ids)
            
            return {
                'id': msg_id,
                'subject': subject,
                'sender_name': sender_name or sender_email.split('@')[0],
                'sender_email': sender_email,
                'to': to_header,
                'date': date_header,
                'timestamp': message.get('internalDate'),
                'snippet': message.get('snippet', ''),
                'body_preview': body[:500] if body else '',
                'labels': label_ids,
                'category': category,
                'is_unread': 'UNREAD' in label_ids,
                'is_starred': 'STARRED' in label_ids,
                'has_attachments': self._has_attachments(message.get('payload', {}))
            }
            
        except HttpError as error:
            print(f"⚠️ Error fetching message {msg_id}: {error}")
            return None
    
    def _get_header(self, headers: List[Dict], name: str) -> Optional[str]:
        """Extract a specific header value from headers list."""
        for header in headers:
            if header['name'].lower() == name.lower():
                return header['value']
        return None
    
    def _get_email_body(self, payload: Dict) -> str:
        """Extract email body text from payload."""
        body = ''
        
        if 'body' in payload and payload['body'].get('data'):
            body = base64.urlsafe_b64decode(payload['body']['data']).decode('utf-8', errors='ignore')
        elif 'parts' in payload:
            for part in payload['parts']:
                if part['mimeType'] == 'text/plain':
                    if part['body'].get('data'):
                        body = base64.urlsafe_b64decode(part['body']['data']).decode('utf-8', errors='ignore')
                        break
                elif part['mimeType'] == 'text/html' and not body:
                    if part['body'].get('data'):
                        body = base64.urlsafe_b64decode(part['body']['data']).decode('utf-8', errors='ignore')
                elif 'parts' in part:
                    # Handle nested multipart
                    body = self._get_email_body(part)
                    if body:
                        break
        
        return body
    
    def _has_attachments(self, payload: Dict) -> bool:
        """Check if email has attachments."""
        if 'parts' in payload:
            for part in payload['parts']:
                if part.get('filename'):
                    return True
                if 'parts' in part:
                    if self._has_attachments(part):
                        return True
        return False
    
    def _categorize_email(self, label_ids: List[str]) -> str:
        """Categorize email based on its labels."""
        if 'SPAM' in label_ids:
            return 'Spam'
        elif 'TRASH' in label_ids:
            return 'Trash'
        elif 'CATEGORY_PROMOTIONS' in label_ids:
            return 'Promotions'
        elif 'CATEGORY_SOCIAL' in label_ids:
            return 'Social'
        elif 'CATEGORY_UPDATES' in label_ids:
            return 'Updates'
        elif 'CATEGORY_FORUMS' in label_ids:
            return 'Forums'
        elif 'CATEGORY_PERSONAL' in label_ids:
            return 'Personal'
        elif 'SENT' in label_ids:
            return 'Sent'
        elif 'DRAFT' in label_ids:
            return 'Draft'
        elif 'INBOX' in label_ids:
            return 'Primary'
        else:
            return 'Other'
    
    def categorize_emails(self, emails: List[Dict]) -> Dict[str, List[Dict]]:
        """
        Organize emails by their category.
        
        Args:
            emails: List of email dictionaries
            
        Returns:
            Dictionary mapping category names to email lists
        """
        categories = {}
        
        for email in emails:
            category = email['category']
            if category not in categories:
                categories[category] = []
            categories[category].append(email)
        
        return categories
    
    def save_to_json(self, emails: List[Dict], filename: str = 'all_emails.json'):
        """Save emails to JSON file."""
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        filepath = os.path.join(OUTPUT_DIR, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(emails, f, indent=2, ensure_ascii=False)
        
        print(f"💾 Saved {len(emails)} emails to {filepath}")
    
    def save_categorized(self, categorized: Dict[str, List[Dict]]):
        """Save categorized emails to separate JSON files."""
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        
        # Save summary
        summary = {cat: len(emails) for cat, emails in categorized.items()}
        summary_path = os.path.join(OUTPUT_DIR, 'summary.json')
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2)
        
        # Save each category
        for category, emails in categorized.items():
            filename = f"{category.lower().replace(' ', '_')}_emails.json"
            filepath = os.path.join(OUTPUT_DIR, filename)
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(emails, f, indent=2, ensure_ascii=False)
            
            print(f"   📁 {category}: {len(emails)} emails -> {filename}")
    
    def save_to_csv(self, emails: List[Dict], filename: str = 'all_emails.csv'):
        """Save emails to CSV file using pandas."""
        import pandas as pd
        
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        filepath = os.path.join(OUTPUT_DIR, filename)
        
        # Prepare data for CSV (exclude complex fields)
        csv_data = []
        for email in emails:
            csv_data.append({
                'Date': email['date'],
                'Category': email['category'],
                'Sender Name': email['sender_name'],
                'Sender Email': email['sender_email'],
                'Subject': email['subject'],
                'Snippet': email['snippet'],
                'Unread': email['is_unread'],
                'Starred': email['is_starred'],
                'Has Attachments': email['has_attachments']
            })
        
        df = pd.DataFrame(csv_data)
        df.to_csv(filepath, index=False, encoding='utf-8')
        
        print(f"📊 Saved to CSV: {filepath}")
    
    def print_summary(self, categorized: Dict[str, List[Dict]]):
        """Print a summary of categorized emails."""
        print("\n" + "="*60)
        print("📊 EMAIL SUMMARY")
        print("="*60)
        
        total = sum(len(emails) for emails in categorized.values())
        
        # Sort categories by count
        sorted_cats = sorted(categorized.items(), key=lambda x: len(x[1]), reverse=True)
        
        for category, emails in sorted_cats:
            bar_length = int((len(emails) / total) * 40) if total > 0 else 0
            bar = "█" * bar_length
            print(f"  {category:15} │ {len(emails):5} │ {bar}")
        
        print("-"*60)
        print(f"  {'TOTAL':15} │ {total:5}")
        print("="*60)


def main():
    """Main function to run the Gmail parser."""
    print("\n" + "="*60)
    print("📧 GMAIL MAIL PARSER")
    print("="*60)
    
    parser = GmailParser()
    
    # Step 1: Authenticate
    if not parser.authenticate():
        return
    
    # Step 2: Fetch all emails
    print("\n📥 Fetching emails from all categories...")
    emails = parser.fetch_emails(max_results=500)
    
    if not emails:
        print("⚠️ No emails found!")
        return
    
    # Step 3: Categorize emails
    print("\n🏷️ Categorizing emails...")
    categorized = parser.categorize_emails(emails)
    
    # Step 4: Print summary
    parser.print_summary(categorized)
    
    # Step 5: Save results
    print("\n💾 Saving parsed emails...")
    parser.save_to_json(emails)
    parser.save_categorized(categorized)
    parser.save_to_csv(emails)
    
    print("\n✅ Done! Check the 'parsed_emails' folder for results.")
    print("="*60 + "\n")


if __name__ == '__main__':
    main()