#!/usr/bin/env python3

import imaplib
import email
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders
import time
import sys

def create_simple_message(subject, body, from_addr="test-sender@example.com", to_addr="user@example.com"):
    """Create a simple text message"""
    msg = MIMEText(body)
    msg['Subject'] = subject
    msg['From'] = from_addr
    msg['To'] = to_addr
    msg['Date'] = email.utils.formatdate(localtime=True)
    return msg.as_string()

def create_message_with_attachment(subject, body, attachment_content, from_addr="test-sender@example.com", to_addr="user@example.com"):
    """Create a message with an attachment"""
    msg = MIMEMultipart()
    msg['Subject'] = subject
    msg['From'] = from_addr
    msg['To'] = to_addr
    msg['Date'] = email.utils.formatdate(localtime=True)
    
    # Add body
    msg.attach(MIMEText(body, 'plain'))
    
    # Add attachment
    part = MIMEBase('text', 'plain')
    part.set_payload(attachment_content)
    encoders.encode_base64(part)
    part.add_header('Content-Disposition', 'attachment; filename="attachment.txt"')
    msg.attach(part)
    
    return msg.as_string()

def populate_user_mailbox(username, password, host='localhost', port=3143):
    """Populate a user's mailbox with test messages"""
    print(f"Connecting to {username}@{host}:{port}")
    
    try:
        # Connect to IMAP server
        imap = imaplib.IMAP4(host, port)
        imap.login(username, password)
        
        # Create additional folders for testing wildcards
        # These folders are designed to test various wildcard patterns
        test_folders = [
            'Sent',                    # Exact match
            'Work/Projects',           # Work/* pattern
            'Work/Archive',            # Work/* pattern  
            'Work/Temp',               # Work/* but excluded by *Temp*
            'Personal',                # Exact match
            'Important/Urgent',        # Important/* pattern
            'Important/Meetings',      # Important/* pattern
            'Archive/2024',            # Archive/* pattern
            'Archive/Projects',        # Archive/* pattern
            'Archive/Drafts',          # Archive/* but excluded by */Drafts
            'Drafts',                  # Should be excluded
            'Trash'                    # Should be excluded
        ]
        created_folders = ['INBOX']  # INBOX always exists
        
        for folder in test_folders:
            try:
                imap.create(folder)
                created_folders.append(folder)
                print(f"  Created folder: {folder}")
            except Exception as e:
                print(f"  Folder {folder} may already exist or creation failed: {e}")
                # Try to select it to see if it exists
                try:
                    imap.select(folder)
                    created_folders.append(folder)
                    print(f"  Folder {folder} already exists")
                except:
                    pass
        
        print(f"Available folders for {username}: {created_folders}")
        
        # Populate each folder with messages
        for folder in created_folders[:3]:  # Limit to first 3 folders to avoid too many messages
            imap.select(folder)
            print(f"Creating messages in {folder} for {username}...")
            
            # Create fewer messages per folder for testing
            message_count = 3 if folder != 'INBOX' else 10
            
            for i in range(1, message_count + 1):
                if i % 3 == 0 and folder == 'INBOX':
                    # Every 3rd message has attachment (only in INBOX)
                    msg = create_message_with_attachment(
                        f"[{folder}] Test Message {i} with Attachment",
                        f"This is test message {i} in {folder} for {username} with an attachment.",
                        f"Sample attachment content for message {i} in {folder}",
                        "test-sender@example.com",
                        f"{username}@example.com"
                    )
                    print(f"    Created message {i} (with attachment)")
                else:
                    # Simple message
                    msg = create_simple_message(
                        f"[{folder}] Test Message {i}",
                        f"This is test message {i} in {folder} for {username}.",
                        "test-sender@example.com",
                        f"{username}@example.com"
                    )
                    print(f"    Created message {i}")
                
                # Append message to current folder
                imap.append(folder, None, None, msg.encode('utf-8'))
                time.sleep(0.1)  # Small delay to avoid overwhelming
            
            # Add special messages only to INBOX for keyword filtering tests
            if folder == 'INBOX':
                special_messages = [
                    ("Important Meeting Reminder", "This is an important meeting message for testing keyword filters.", "manager@company.com", "team@company.com"),
                    ("Urgent: System Maintenance", "Important notification about system maintenance.", "admin@company.com", f"{username}@example.com"),
                    ("Regular Newsletter", "This is a regular newsletter message.", "newsletter@external.com", f"{username}@example.com"),
                    ("Team Meeting Notes", "Meeting notes from the team.", "secretary@company.com", "support@company.com"),
                    ("Message with Special Characters", "This message contains special characters: äöü ñ 中文 🚀", "test-sender@example.com", f"{username}@example.com")
                ]
                
                for subject, body, sender, recipient in special_messages:
                    msg = create_simple_message(subject, body, sender, recipient)
                    imap.append(folder, None, None, msg.encode('utf-8'))
                    print(f"    Created special message: {subject} from {sender} to {recipient}")
                    time.sleep(0.1)
        
        imap.logout()
        print(f"✅ Successfully created messages across {len(created_folders[:3])} folders for {username}")
        return True
        
    except Exception as e:
        print(f"❌ Error populating {username}: {e}")
        return False

def main():
    print("🚀 Populating GreenMail with test messages using IMAP...")
    
    # Test accounts
    accounts = [
        ("testuser1", "password123"),
        ("testuser2", "password456"), 
        ("syncuser", "syncpass"),
        ("archiveuser", "archivepass")
    ]
    
    # Wait for GreenMail to be ready
    print("Waiting for GreenMail to be ready...")
    time.sleep(5)
    
    success_count = 0
    for username, password in accounts:
        if populate_user_mailbox(username, password):
            success_count += 1
        time.sleep(1)  # Brief pause between accounts
    
    print(f"\n🎉 Successfully populated {success_count}/{len(accounts)} accounts!")
    
    if success_count == len(accounts):
        print("\n✅ All test accounts ready:")
        for username, password in accounts:
            print(f"  - {username}:{password}@localhost")
        print(f"\nGreenMail Services:")
        print(f"  - IMAP: localhost:3143")
        print(f"  - IMAPS: localhost:3993") 
        print(f"  - SMTP: localhost:3025")
        return 0
    else:
        print(f"\n❌ Failed to populate {len(accounts) - success_count} accounts")
        return 1

if __name__ == "__main__":
    sys.exit(main())