"""
Gmail Cleanup Agent v2
======================
A generic, interactive Gmail cleanup agent.
Safely preview and delete emails matching any search query.

SETUP:
------
1. pip install google-auth-oauthlib google-auth-httplib2 google-api-python-client
2. Place your credentials.json in the same folder
3. Run: python3 gmail_cleanup_agent_v2.py

FIRST TIME: A browser will open for Google authentication.
Subsequent runs use saved token — no browser needed.
"""

import time
import os
import pickle
import argparse
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
from googleapiclient.discovery import build

# ─── CONSTANTS ────────────────────────────────────────────────────────────────

SCOPES = ["https://mail.google.com/"]
CREDENTIALS_FILE = "credentials.json"
TOKEN_FILE = "token.pickle"
BATCH_SIZE = 100

# ─── COLORS FOR TERMINAL OUTPUT ───────────────────────────────────────────────

class Color:
    GREEN  = "\033[92m"
    YELLOW = "\033[93m"
    RED    = "\033[91m"
    BLUE   = "\033[94m"
    BOLD   = "\033[1m"
    RESET  = "\033[0m"

def green(text):  return f"{Color.GREEN}{text}{Color.RESET}"
def yellow(text): return f"{Color.YELLOW}{text}{Color.RESET}"
def red(text):    return f"{Color.RED}{text}{Color.RESET}"
def blue(text):   return f"{Color.BLUE}{text}{Color.RESET}"
def bold(text):   return f"{Color.BOLD}{text}{Color.RESET}"

# ─── AUTH ─────────────────────────────────────────────────────────────────────

def authenticate():
    """Authenticate with Gmail API. Opens browser on first run."""
    creds = None

    if os.path.exists(TOKEN_FILE):
        with open(TOKEN_FILE, "rb") as token:
            creds = pickle.load(token)

    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            print(yellow("🔄 Refreshing authentication token..."))
            creds.refresh(Request())
        else:
            if not os.path.exists(CREDENTIALS_FILE):
                print(red(f"❌ '{CREDENTIALS_FILE}' not found!"))
                print(yellow("   Download it from Google Cloud Console → APIs & Services → Credentials"))
                exit(1)
            print(yellow("🌐 Opening browser for authentication..."))
            flow = InstalledAppFlow.from_client_secrets_file(CREDENTIALS_FILE, SCOPES)
            creds = flow.run_local_server(port=0)

        with open(TOKEN_FILE, "wb") as token:
            pickle.dump(creds, token)

    return build("gmail", "v1", credentials=creds)

# ─── GMAIL HELPERS ────────────────────────────────────────────────────────────

def search_emails(service, query, max_results=BATCH_SIZE):
    """Search emails matching the query."""
    result = service.users().messages().list(
        userId="me",
        q=query,
        maxResults=max_results
    ).execute()
    return result.get("messages", [])


def get_email_summary(service, msg_id):
    """Get a short summary (sender + subject) of an email."""
    try:
        msg = service.users().messages().get(
            userId="me", id=msg_id, format="metadata",
            metadataHeaders=["From", "Subject"]
        ).execute()
        headers = {h["name"]: h["value"] for h in msg["payload"]["headers"]}
        sender  = headers.get("From", "Unknown")[:40]
        subject = headers.get("Subject", "(no subject)")[:60]
        return f"{sender:<42} | {subject}"
    except Exception:
        return f"(could not fetch details for {msg_id})"


def trash_email(service, msg_id):
    """Move email to Trash (recoverable for 30 days)."""
    service.users().messages().trash(userId="me", id=msg_id).execute()


def delete_email(service, msg_id):
    """Permanently delete an email. Cannot be undone."""
    service.users().messages().delete(userId="me", id=msg_id).execute()

# ─── DRY RUN PREVIEW ──────────────────────────────────────────────────────────

def dry_run_preview(service, query):
    """Preview up to 20 matching emails without deleting anything."""
    print(blue("\n📋 DRY RUN — Previewing matching emails (no deletions)"))
    print(blue("─" * 70))

    messages = search_emails(service, query, max_results=20)

    if not messages:
        print(green("✅ No emails found matching this query."))
        return 0

    # Count total matches
    count_result = service.users().messages().list(
        userId="me", q=query, maxResults=1
    ).execute()
    total = count_result.get("resultSizeEstimate", len(messages))

    print(f"\n{bold('Estimated total matches:')} {yellow(str(total))} {yellow('(⚠️  Gmail API caps this estimate at ~201 — real count may be much higher)')}")
    print(f"{bold('Showing first:')} {len(messages)}\n")
    print(f"{'SENDER':<42} | {'SUBJECT'}")
    print("─" * 70)

    for msg in messages:
        print(get_email_summary(service, msg["id"]))

    if total > 20:
        print(f"\n  ... and approximately {total - 20} more emails")

    return total

# ─── MAIN AGENT LOOP ──────────────────────────────────────────────────────────

def run_cleanup(service, query, mode):
    """Main agent loop — delete emails until query returns zero results."""
    print(blue(f"\n🤖 Starting cleanup in {bold(mode.upper())} mode"))
    print(blue("─" * 70))

    total_deleted = 0
    round_number  = 0

    while True:
        round_number += 1
        print(yellow(f"\n🔍 Round {round_number}: Searching for matching emails..."))

        messages = search_emails(service, query)

        if not messages:
            print(green(f"\n✅ Done! No more emails matching the query."))
            print(green(f"📊 Total emails cleaned up: {bold(str(total_deleted))}"))
            break

        print(f"   Found {yellow(str(len(messages)))} emails — processing...")

        for i, msg in enumerate(messages):
            try:
                if mode == "trash":
                    trash_email(service, msg["id"])
                else:
                    delete_email(service, msg["id"])

                total_deleted += 1

                if (i + 1) % 10 == 0:
                    print(f"   ... processed {i + 1}/{len(messages)}")

            except Exception as e:
                print(red(f"   ⚠️  Failed on message {msg['id']}: {e}"))

        print(green(f"   ✅ Round {round_number} complete — {len(messages)} emails removed"))
        time.sleep(0.5)  # Avoid Gmail API rate limits

    return total_deleted

# ─── INTERACTIVE MODE ─────────────────────────────────────────────────────────

def interactive_mode(service):
    """Walk the user through query input and confirmation interactively."""
    print(bold(blue("\n╔══════════════════════════════════════╗")))
    print(bold(blue("║      Gmail Cleanup Agent v2 🤖       ║")))
    print(bold(blue("╚══════════════════════════════════════╝")))

    print("""
Example queries:
  from:linkedin.com                        → All LinkedIn emails
  from:linkedin.com OR from:facebook.com   → LinkedIn + Facebook
  is:unread older_than:1y                  → Unread emails older than 1 year
  category:promotions older_than:30d       → Promotions older than 30 days
  subject:newsletter                       → Emails with "newsletter" in subject
""")

    # Get query
    query = input(bold("Enter your Gmail search query:\n> ")).strip()
    if not query:
        print(red("❌ No query entered. Exiting."))
        return

    # Dry run first
    total = dry_run_preview(service, query)

    if total == 0:
        return

    # Ask for action
    print(f"""
{bold("What would you like to do?")}
  1. Trash emails    (recoverable from Trash for 30 days) ✅ recommended
  2. Delete forever  (permanent, cannot be undone) ⚠️
  3. Cancel
""")
    choice = input(bold("Enter choice (1/2/3): ")).strip()

    if choice == "1":
        mode = "trash"
    elif choice == "2":
        confirm = input(red(f"⚠️  Are you sure you want to PERMANENTLY delete ~{total} emails? (yes/no): ")).strip().lower()
        if confirm != "yes":
            print(yellow("Cancelled."))
            return
        mode = "delete"
    else:
        print(yellow("Cancelled."))
        return

    # Run cleanup
    run_cleanup(service, query, mode)

# ─── CLI MODE ─────────────────────────────────────────────────────────────────

def parse_args():
    parser = argparse.ArgumentParser(
        description="Gmail Cleanup Agent — delete emails matching a search query",
        formatter_class=argparse.RawTextHelpFormatter
    )
    parser.add_argument(
        "--query", "-q",
        type=str,
        help='Gmail search query e.g. "from:linkedin.com OR from:facebook.com"'
    )
    parser.add_argument(
        "--mode", "-m",
        choices=["trash", "delete"],
        default="trash",
        help="trash = move to Trash (default, recoverable)\ndelete = permanent delete"
    )
    parser.add_argument(
        "--dry-run", "-d",
        action="store_true",
        help="Preview matching emails without deleting anything"
    )
    return parser.parse_args()

# ─── ENTRY POINT ──────────────────────────────────────────────────────────────

def main():
    args = parse_args()

    print(yellow("🔐 Authenticating with Gmail..."))
    service = authenticate()
    print(green("✅ Authenticated!\n"))

    if args.query:
        # Non-interactive CLI mode
        if args.dry_run:
            dry_run_preview(service, args.query)
        else:
            run_cleanup(service, args.query, args.mode)
    else:
        # Interactive mode
        interactive_mode(service)


if __name__ == "__main__":
    main()
