import msal
import imaplib
import os
import csv
import re
import time
import socket

# 🚨 Force Python to drop the connection if Microsoft stops responding (tarpitting)
socket.setdefaulttimeout(30)

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

AUTHORITY = "https://login.microsoftonline.com/884a3c53-8a5a-4d79-b0e0-a62ab5a794a1"
# MSAL automatically requests offline_access, so we only list the IMAP scope here
SCOPES = ["https://outlook.office.com/IMAP.AccessAsUser.All"]

# ==========================================
# 2. HELPER FUNCTION: GENERATE CSV
# ==========================================
def generate_user_csv(email_address, first_name="Guillem", last_name="Hernandez Sola"):
    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}")

# ==========================================
# 3. TOKEN & CONNECTION MANAGERS
# ==========================================
def get_valid_token(app):
    """Checks the cache for a valid token, refreshes silently if needed, or prompts user."""
    accounts = app.get_accounts()
    if accounts:
        result = app.acquire_token_silent(SCOPES, account=accounts[0])
        if result and "access_token" in result:
            return result["access_token"]

    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("⏳ Waiting for browser authentication...")
    
    result = app.acquire_token_by_device_flow(flow)
    if "access_token" not in result:
        raise Exception(f"Failed to get token: {result.get('error_description')}")
        
    return result["access_token"]

def connect_to_imap(email, token):
    """Creates a fresh, authenticated connection to the IMAP server."""
    auth_string = f"user={email}\x01auth=Bearer {token}\x01\x01"
    mail = imaplib.IMAP4_SSL("outlook.office365.com", 993)
    mail.authenticate("XOAUTH2", lambda x: auth_string.encode("utf-8"))
    return mail

# ==========================================
# 4. MAIN EXECUTION
# ==========================================
def main():
    print("🔄 Initializing Microsoft Authentication...")
    app = msal.PublicClientApplication(CLIENT_ID, authority=AUTHORITY)
    
    try:
        access_token = get_valid_token(app)
        print("\n✅ Access Token Acquired Successfully!")
        
        print("\n🔌 Connecting to Outlook IMAP server...")
        mail = connect_to_imap(EMAIL_ADDRESS, access_token)
        print("✅ Successfully logged into IMAP via OAuth2!")
    except Exception as e:
        print(f"\n❌ Authentication failed: {e}")
        return

    try:
        username = EMAIL_ADDRESS.split("@")[0]
        base_download_dir = f"downloaded_emails_{username}"
        os.makedirs(base_download_dir, exist_ok=True)

        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')
                
                if "\\Noselect" in folder_string:
                    continue
                
                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}")
                
                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...")
                
                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 loop with Reconnect & Timeout Logic
                for e_id in email_ids:
                    file_path = os.path.join(folder_dir, f"email_{e_id.decode('utf-8')}.eml")
                    
                    # 🚀 SKIP EXISTING: Don't re-download emails we already have!
                    if os.path.exists(file_path):
                        continue
                    
                    success = False
                    while not success:
                        try:
                            status, msg_data = mail.fetch(e_id, "(RFC822)")
                            for response_part in msg_data:
                                if isinstance(response_part, tuple):
                                    raw_email = response_part[1]
                                    with open(file_path, "wb") as f:
                                        f.write(raw_email)
                            success = True
                            
                        # Catch token expiration, forced disconnects, AND silent timeouts
                        except (imaplib.IMAP4.abort, imaplib.IMAP4.error, ConnectionResetError, socket.timeout, TimeoutError) as e:
                            print(f"\n   ⚠️ Connection lost or timed out. Refreshing token and reconnecting...")
                            try:
                                access_token = get_valid_token(app)
                                mail = connect_to_imap(EMAIL_ADDRESS, access_token)
                                mail.select(f'"{folder_name}"', readonly=True) 
                                print("   ✅ Reconnected! Resuming download...")
                            except Exception as reconnect_error:
                                print(f"   ❌ Reconnection failed: {reconnect_error}. Retrying in 5 seconds...")
                                time.sleep(5)
                                
            print(f"\n🎉 All folders successfully downloaded to '{base_download_dir}'!")
        
        mail.logout()
        print("👋 Logged out successfully.\n")

        print("⚙️ Generating configuration files...")
        generate_user_csv(EMAIL_ADDRESS)
        print("🎉 Migration prep complete!")

    except Exception as e:
        print(f"\n❌ A critical error occurred: {e}")

if __name__ == "__main__":
    main()