mcp-maildir/src/indexer.py

"""
Indexer script to parse emails from Maildir and push them to Qdrant.
"""

import os
import email
import mailbox
import warnings
from concurrent.futures import ThreadPoolExecutor, Future
from datetime import datetime
from email.utils import parsedate_to_datetime, parseaddr
from email.header import decode_header
from typing import List, Dict, Any, Tuple
import uuid

from dotenv import load_dotenv
from qdrant_client import QdrantClient
from qdrant_client.http import models
from fastembed import TextEmbedding
from bs4 import BeautifulSoup

# Load .env config
load_dotenv()

# Configuration
MAILDIR_PATH = os.environ.get("MAILDIR_PATH", "")
MAILDIR_FOLDERS = os.environ.get("MAILDIR_FOLDERS", "")
QDRANT_URL = os.environ.get("QDRANT_URL", "")
COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "")

if not MAILDIR_PATH:
    raise ValueError("MAILDIR_PATH environment variable is required.")
if not QDRANT_URL:
    raise ValueError("QDRANT_URL environment variable is required.")
if not COLLECTION_NAME:
    raise ValueError("COLLECTION_NAME environment variable is required.")

EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME", "BAAI/bge-small-en-v1.5")
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "100"))
EMBEDDING_BATCH_SIZE = int(os.environ.get("EMBEDDING_BATCH_SIZE", "64"))
METADATA_COLLECTION = "mcp_indexer_metadata"
INCREMENTAL_DAYS = int(os.environ.get("INCREMENTAL_DAYS", "7"))
FORCE_REINDEX = os.environ.get("FORCE_REINDEX", "").lower() in ("1", "true", "yes")

def decode_mime_words(s: str) -> str:
    """Decodes MIME encoded strings (e.g. subjects, filenames)."""
    if not s:
        return ""
    decoded_words = decode_header(s)
    result = []
    for word, encoding in decoded_words:
        if isinstance(word, bytes):
            try:
                result.append(word.decode(encoding or 'utf-8', errors='replace'))
            except LookupError:
                result.append(word.decode('utf-8', errors='replace'))
        else:
            result.append(word)
    return "".join(result)

def extract_text_from_html(html_content: str) -> str:
    """Extracts plain text from HTML content."""
    try:
        soup = BeautifulSoup(html_content, "html.parser")
        return soup.get_text(separator=" ", strip=True)
    except Exception:
        return html_content

def normalize_email_address(value: str) -> str:
    """Extracts and normalizes the bare email address from a header value."""
    if not value:
        return ""
    _, addr = parseaddr(value)
    return (addr or value).strip().lower()


def parse_email_message(msg: mailbox.Message, email_id: str = "") -> Tuple[str, List[str]]:
    """Extracts plain text body and a list of attachment filenames."""
    body_parts = []
    attachments = []

    for part in msg.walk():
        # Skip multiparts, we only care about leaf nodes
        if part.is_multipart():
            continue
        
        content_type = part.get_content_type()
        content_disposition = str(part.get("Content-Disposition", ""))

        # Check for attachments
        if "attachment" in content_disposition or part.get_filename():
            filename = part.get_filename()
            if filename:
                attachments.append(decode_mime_words(filename))
            continue
            
        # Extract text body
        if content_type in ["text/plain", "text/html"]:
            try:
                payload = part.get_payload(decode=True)
                if payload:
                    charset = part.get_content_charset('utf-8') or 'utf-8'
                    if isinstance(payload, bytes):
                        try:
                            text = payload.decode(charset, errors='replace')
                        except (LookupError, UnicodeDecodeError):
                            # Unknown or broken charset — fall back to utf-8
                            print(f"  Warning: unknown charset '{charset}', falling back to utf-8 [{email_id}]")
                            text = payload.decode('utf-8', errors='replace')
                    else:
                        text = str(payload)
                        
                    if content_type == "text/html":
                        with warnings.catch_warnings(record=True) as caught:
                            warnings.simplefilter("always")
                            text = extract_text_from_html(text)
                            for w in caught:
                                print(f"  Warning: {w.category.__name__}: {w.message} [{email_id}]")
                    body_parts.append(text)
            except Exception as e:
                print(f"  Warning: error extracting payload: {e} [{email_id}]")

    return "\n".join(body_parts).strip(), attachments

def init_qdrant_collection(client: QdrantClient, vector_size: int):
    """Ensures Qdrant collection exists and payload indexes are created."""
    
    # Check if collection exists
    collections = client.get_collections().collections
    if not any(c.name == COLLECTION_NAME for c in collections):
        print(f"Creating collection '{COLLECTION_NAME}' with vector size {vector_size}...")
        client.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE),
        )
    else:
        print(f"Collection '{COLLECTION_NAME}' already exists.")

    # Create payload indexes for filtering metadata deterministically
    print("Ensuring payload indexes exist...")
    
    # Date index (DATETIME)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="date",
        field_schema=models.PayloadSchemaType.DATETIME,
    )
    
    # Sender index (KEYWORD)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="sender",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )
    
    # Receiver index (KEYWORD)
    client.create_payload_index(
        collection_name=COLLECTION_NAME,
        field_name="receiver",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )

def init_metadata_collection(client: QdrantClient):
    """Ensures the indexer metadata collection exists in Qdrant."""
    collections = client.get_collections().collections
    if not any(c.name == METADATA_COLLECTION for c in collections):
        print(f"Creating metadata collection '{METADATA_COLLECTION}'...")
        client.create_collection(
            collection_name=METADATA_COLLECTION,
            # Minimal vector (size=1) — we only use this collection for payload storage
            vectors_config=models.VectorParams(size=1, distance=models.Distance.COSINE),
        )


def is_bootstrap_done(client: QdrantClient) -> bool:
    """Returns True if a successful full bootstrap has already been recorded."""
    try:
        results, _ = client.scroll(
            collection_name=METADATA_COLLECTION,
            scroll_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="event",
                        match=models.MatchValue(value="bootstrap_complete"),
                    )
                ]
            ),
            limit=1,
        )
        return len(results) > 0
    except Exception as e:
        print(f"Warning: could not check bootstrap state: {e}")
        return False


def mark_bootstrap_done(client: QdrantClient):
    """Records a bootstrap_complete event in the metadata collection."""
    point_id = str(uuid.uuid5(uuid.NAMESPACE_OID, "bootstrap_complete"))
    client.upsert(
        collection_name=METADATA_COLLECTION,
        points=[
            models.PointStruct(
                id=point_id,
                vector=[0.0],  # placeholder — collection is payload-only
                payload={
                    "event": "bootstrap_complete",
                    "timestamp": datetime.now().isoformat(),
                },
            )
        ],
    )
    print("Bootstrap state recorded in Qdrant metadata collection.")


def get_recent_keys(mbox: mailbox.Maildir, days: int) -> set:
    """
    Returns the set of Maildir keys whose backing file has been modified
    within the last `days` days (based on filesystem mtime).
    """
    from datetime import timezone, timedelta

    cutoff = datetime.now(tz=timezone.utc) - timedelta(days=days)
    cutoff_ts = cutoff.timestamp()
    recent = set()
    maildir_root = mbox._path  # type: ignore[attr-defined]

    for subdir in ("cur", "new"):
        subdir_path = os.path.join(maildir_root, subdir)
        try:
            filenames = os.listdir(subdir_path)
        except OSError:
            continue

        for filename in filenames:
            file_path = os.path.join(subdir_path, filename)
            try:
                if os.stat(file_path).st_mtime >= cutoff_ts:
                    # Maildir keys are the filename, but the actual file on disk
                    # often has a suffix (like ":2,S"). mailbox.Maildir treats
                    # the part BEFORE the first colon as the key.
                    key = filename.split(":")[0] if ":" in filename else filename
                    recent.add(key)
            except OSError:
                continue

    return recent


def main():
    """
    Main ingestion function.
    Reads Maildir, extracts text, generates local embeddings, and pushes to Qdrant.
    """
    print(f"Indexing emails from {MAILDIR_PATH} into {QDRANT_URL}...")

    if not os.path.exists(MAILDIR_PATH):
        print(f"Error: Maildir path not found: {MAILDIR_PATH}")
        return

    # Initialize model
    print(f"Loading embedding model: {EMBEDDING_MODEL_NAME}...")
    model = TextEmbedding(model_name=EMBEDDING_MODEL_NAME)
    vector_size = len(next(iter(model.embed(["dimension_probe"]))))

    # Initialize Qdrant
    print("Connecting to Qdrant...")
    qdrant_client = QdrantClient(url=QDRANT_URL)

    # Force reindex: wipe existing collections to start from scratch
    if FORCE_REINDEX:
        print("[FORCE_REINDEX] Deleting existing collections for a clean re-bootstrap...")
        for col_name in (COLLECTION_NAME, METADATA_COLLECTION):
            try:
                qdrant_client.delete_collection(collection_name=col_name)
                print(f"  Deleted collection '{col_name}'.")
            except Exception:
                print(f"  Collection '{col_name}' did not exist, skipping.")

    init_qdrant_collection(qdrant_client, vector_size)
    init_metadata_collection(qdrant_client)

    # Determine indexing mode: full bootstrap or incremental update
    if is_bootstrap_done(qdrant_client):
        mode = "incremental"
        print(f"[MODE] INCREMENTAL — scanning only files modified in the last {INCREMENTAL_DAYS} days.")
    else:
        mode = "full"
        print("[MODE] BOOTSTRAP — full scan of all emails (first-time indexing).")

    # Pipeline:
    #   parse phase  → accumulates (vector_text, payload, point_id) into a batch
    #   embed phase  → model.embed() called once per batch (vectorized ONNX inference)
    #   upsert phase → submitted to a background thread so the next batch can be
    #                  parsed+embedded while the previous one is in-flight to Qdrant
    #
    # pending_batch: accumulates parsed email metadata + vector_text until BATCH_SIZE
    # pending_future: the in-flight ThreadPoolExecutor Future for the previous upsert
    pending_batch: List[Dict[str, Any]] = []
    pending_future: Future | None = None
    has_error = False
    total_processed = 0

    def _flush_batch(executor: ThreadPoolExecutor, batch: List[Dict[str, Any]]) -> Future:
        """Embed a batch of pre-parsed emails and submit an async upsert to Qdrant."""
        vector_texts = [item["vector_text"] for item in batch]
        vectors = [v.tolist() for v in model.embed(vector_texts, batch_size=EMBEDDING_BATCH_SIZE)]
        points = [
            models.PointStruct(
                id=item["point_id"],
                vector=vectors[i],
                payload=item["payload"],
            )
            for i, item in enumerate(batch)
        ]
        return executor.submit(
            qdrant_client.upsert,
            collection_name=COLLECTION_NAME,
            points=points,
        )

    with ThreadPoolExecutor(max_workers=1) as executor:
        # Determine which directories to process
        maildir_roots = []
        if MAILDIR_FOLDERS:
            folders = [f.strip() for f in MAILDIR_FOLDERS.split(",") if f.strip()]
            for folder in folders:
                folder_path = os.path.join(MAILDIR_PATH, folder)
                if os.path.isdir(folder_path):
                    maildir_roots.append(folder_path)
                else:
                    print(f"Warning: Specified folder not found: {folder_path}")
        else:
            for root_dir, dirs, _ in os.walk(MAILDIR_PATH):
                if all(subdir in dirs for subdir in ['cur', 'new', 'tmp']):
                    maildir_roots.append(root_dir)

        # Iterate and parse over selected maildir directories
        for root in maildir_roots:
            # A valid Maildir has 'cur', 'new', and 'tmp' subdirectories
            if not all(os.path.isdir(os.path.join(root, subdir)) for subdir in ['cur', 'new', 'tmp']):
                continue

            print(f"Processing Maildir found at: {root}")
            mbox = mailbox.Maildir(root)
            total_emails_in_dir = len(mbox)

            if mode == "incremental":
                keys_to_process = get_recent_keys(mbox, INCREMENTAL_DAYS)
                print(
                    f"Found {total_emails_in_dir} emails total, "
                    f"{len(keys_to_process)} modified in the last {INCREMENTAL_DAYS} days."
                )
            else:
                keys_to_process = set(mbox.keys())
                print(f"Found {total_emails_in_dir} emails — indexing all.")

            for key in keys_to_process:
                try:
                    msg = mbox[key]

                    # Parse headers
                    subject = decode_mime_words(msg.get("Subject", "No Subject"))
                    sender_raw = decode_mime_words(msg.get("From", "Unknown"))
                    receiver_raw = decode_mime_words(msg.get("To", "Unknown"))
                    sender = normalize_email_address(sender_raw)
                    receiver = normalize_email_address(receiver_raw)
                    message_id_raw = msg.get("Message-ID")
                    message_id = str(message_id_raw) if message_id_raw is not None else str(uuid.uuid4())

                    # Parse date — msg.get() may return an email.header.Header
                    # object instead of str when the header contains non-ASCII
                    # bytes (e.g. timezone comments like "heure d'été").
                    # We must coerce to str before parsing.
                    date_raw = msg.get("Date")
                    date_str = str(date_raw) if date_raw is not None else None
                    dt_obj = None
                    if date_str:
                        try:
                            dt_obj = parsedate_to_datetime(date_str)
                        except Exception:
                            pass
                    if dt_obj is None:
                        # Fallback: warn and use current time
                        print(f"  Warning: could not parse Date header: {repr(date_raw)} [key={key}, subject={subject}]")
                        dt_obj = datetime.now()
                    iso_date = dt_obj.isoformat()

                    # Parse body and attachments
                    body_text, attachments = parse_email_message(msg, email_id=f"key={key}, subject={subject}")

                    attachments_str = ", ".join(attachments) if attachments else "None"
                    vector_text = (
                        f"Date: {iso_date}\n"
                        f"From: {sender}\n"
                        f"To: {receiver}\n"
                        f"Subject: {subject}\n\n"
                        f"{body_text}\n\n"
                        f"Attachments: {attachments_str}"
                    )

                    pending_batch.append({
                        "vector_text": vector_text,
                        "point_id": str(uuid.uuid5(uuid.NAMESPACE_OID, message_id)),
                        "payload": {
                            "message_id": message_id,
                            "date": iso_date,
                            "sender": sender,
                            "sender_raw": sender_raw,
                            "receiver": receiver,
                            "receiver_raw": receiver_raw,
                            "subject": subject,
                            "body_text": body_text,
                            "attachments": attachments,
                        },
                    })

                    if len(pending_batch) >= BATCH_SIZE:
                        # Wait for the previous upsert to complete before submitting the next
                        if pending_future is not None:
                            pending_future.result()
                        pending_future = _flush_batch(executor, pending_batch)
                        total_processed += len(pending_batch)
                        print(f"  Embedded+upserted batch — {total_processed} emails total so far.")
                        pending_batch = []

                except Exception as e:
                    print(f"Error processing email key={key}: {e}")
                    has_error = True

        # Flush remaining emails
        if pending_batch:
            if pending_future is not None:
                pending_future.result()
            pending_future = _flush_batch(executor, pending_batch)
            total_processed += len(pending_batch)
            pending_batch = []

        # Wait for the last upsert to finish before exiting the executor context
        if pending_future is not None:
            pending_future.result()

    print(f"  Total emails indexed: {total_processed}")

    # Record bootstrap completion so subsequent runs use incremental mode
    if mode == "full" and not has_error:
        mark_bootstrap_done(qdrant_client)

    print("Indexing completed successfully!")

if __name__ == "__main__":
    main()