import msal
import imaplib
import os
import csv
import re  # Added for parsing folder names safely

# ==========================================
# 1. CONFIGURATION
# ==========================================
CLIENT_ID = "05332268-8149-449f-a1f8-1efadd17166f"  
EMAIL_ADDRESS = "marti@agile611.com"  

AUTHORITY = "https://login.microsoftonline.com/884a3c53-8a5a-4d79-b0e0-a62ab5a794a1"
SCOPES = ["https://outlook.office.com/IMAP.AccessAsUser.All"]

# ==========================================
# 2. HELPER FUNCTION: GENERATE CSV
# ==========================================
def generate_user_csv(email_address, first_name="Marti", last_name="Montfort Ruiz"):
    username = email_address.split("@")[0]
    headers = [
        "originUsername", "targetUsername", "password", "pop3enabled", 
        "pop3password", "aliases", "forwards", "filters", 
        "forename", "surname", "mailboxStatus"
    ]
    row_data = {
        "originUsername": username, "targetUsername": username,
        "password": "TempWebmailPassword123!", "pop3enabled": "true",
        "pop3password": "", "aliases": "", "forwards": "", "filters": "",
        "forename": first_name, "surname": last_name, "mailboxStatus": "premium"
    }
    csv_filename = f"{username}_import.csv"
    with open(csv_filename, mode="w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(file, fieldnames=headers, delimiter=";")
        writer.writeheader()
        writer.writerow(row_data)
    print(f"📝 Migration CSV generated: {csv_filename}")

def main():
    # ==========================================
    # 3. GETTING THE OAUTH2 TOKEN
    # ==========================================
    print("🔄 Initializing Microsoft Authentication...")
    app = msal.PublicClientApplication(CLIENT_ID, authority=AUTHORITY)
    flow = app.initiate_device_flow(scopes=SCOPES)
    
    if "user_code" not in flow:
        raise ValueError("Failed to create device flow. Check your Client ID and Azure settings.")

    print("\n🚨 ACTION REQUIRED 🚨")
    print(flow["message"])
    print("\n⏳ Waiting for browser authentication...")
    
    result = app.acquire_token_by_device_flow(flow)
    if "access_token" not in result:
        print("\n❌ Failed to get token:", result.get("error_description"))
        return

    access_token = result["access_token"]
    print("\n✅ Access Token Acquired Successfully!")

    # ==========================================
    # 4. CONNECTING TO IMAP & DOWNLOADING ALL FOLDERS
    # ==========================================
    print("\n🔌 Connecting to Outlook IMAP server...")
    auth_string = f"user={EMAIL_ADDRESS}\x01auth=Bearer {access_token}\x01\x01"

    try:
        mail = imaplib.IMAP4_SSL("outlook.office365.com", 993)
        mail.authenticate("XOAUTH2", lambda x: auth_string.encode("utf-8"))
        print("✅ Successfully logged into IMAP via OAuth2!")
        
        # Base directory for this user
        username = EMAIL_ADDRESS.split("@")[0]
        base_download_dir = f"downloaded_emails_{username}"
        os.makedirs(base_download_dir, exist_ok=True)

        # Fetch all folders in the mailbox
        status, folders = mail.list()
        if status == "OK":
            print(f"📂 Found {len(folders)} folders. Starting full account backup...\n")
            
            for folder_data in folders:
                folder_string = folder_data.decode('utf-8')
                
                # Skip unselectable folders (like root directory markers)
                if "\\Noselect" in folder_string:
                    continue
                
                # Safely extract the folder name (handles spaces and quotes)
                match = re.search(r'\"([^\"]+)\"$', folder_string)
                folder_name = match.group(1) if match else folder_string.split()[-1].strip('"')
                
                print(f"📁 Scanning folder: {folder_name}")
                
                # Select the folder (readonly to prevent accidental modifications)
                status, _ = mail.select(f'"{folder_name}"', readonly=True)
                if status != "OK":
                    print(f"   ⚠️ Could not open {folder_name}. Skipping.")
                    continue
                
                status, data = mail.search(None, "ALL")
                email_ids = data[0].split()
                
                if not email_ids:
                    print("   ↳ Folder is empty.")
                    continue
                
                print(f"   ↳ Found {len(email_ids)} emails. Downloading...")
                
                # Create a safe subfolder name for the OS
                safe_folder_name = "".join([c for c in folder_name if c.isalnum() or c in (' ', '-', '_')]).strip()
                folder_dir = os.path.join(base_download_dir, safe_folder_name)
                os.makedirs(folder_dir, exist_ok=True)
                
                # Download each email into its respective folder
                for e_id in email_ids:
                    status, msg_data = mail.fetch(e_id, "(RFC822)")
                    for response_part in msg_data:
                        if isinstance(response_part, tuple):
                            raw_email = response_part[1]
                            file_path = os.path.join(folder_dir, f"email_{e_id.decode('utf-8')}.eml")
                            with open(file_path, "wb") as f:
                                f.write(raw_email)
                                
            print(f"\n🎉 All folders successfully downloaded to '{base_download_dir}'!")
        
        mail.logout()
        print("👋 Logged out successfully.\n")

        # ==========================================
        # 5. GENERATE THE CSV
        # ==========================================
        print("⚙️ Generating configuration files...")
        generate_user_csv(EMAIL_ADDRESS)
        print("🎉 Migration prep complete!")

    except imaplib.IMAP4.error as e:
        print(f"\n❌ IMAP Authentication failed: {e}")

if __name__ == "__main__":
    main()