"""
PST Email Extractor
====================
Extracts email metadata from Outlook PST files into CSV format.
Supports batch processing of multiple PST files in a directory.

Dependencies:
    pip install tqdm libpff-python
    (see README.md for platform-specific libpff/pypff installation)

Usage:
    python Extract.py [directory]
    If no directory is given, the script's own directory is used.
"""

import pypff
import csv
import sys
import os
import glob
from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor, as_completed
from functools import lru_cache
from tqdm import tqdm


# ─────────────────────────── helpers ────────────────────────────

def _safe_str(value, max_len: int = 0) -> str:
    """Convert a value to a clean string, optionally truncating it."""
    try:
        result = "" if value is None else str(value)
    except Exception:
        return ""
    return result[:max_len] if max_len else result


def get_folder_path(folder) -> str:
    """
    Walk up the parent chain to reconstruct the full folder path.
    Capped at 50 levels to prevent runaway recursion on corrupt files.
    """
    parts: list[str] = []
    current = folder
    for _ in range(50):
        if current is None:
            break
        try:
            name = getattr(current, "name", None)
            parts.insert(0, str(name) if name else "Unknown_Folder")
            current = getattr(current, "parent", None)
        except Exception:
            break
    clean = [p for p in parts if p]
    return "\\".join(clean) if clean else "Root"


def _format_dt(dt_obj) -> str:
    """Return an ISO-8601 string for a datetime-like object, UTC-aware."""
    if dt_obj is None:
        return ""
    try:
        # pypff delivery/submit times are naive UTC – make them explicit
        if dt_obj.tzinfo is None:
            dt_obj = dt_obj.replace(tzinfo=timezone.utc)
        return dt_obj.strftime("%Y-%m-%d %H:%M:%S UTC")
    except Exception:
        return ""


# ──────────────────────── email extraction ──────────────────────

def extract_email_info(message, folder_path: str) -> dict:
    """
    Pull all useful metadata from a pypff message object.
    Every field is individually guarded so one bad attribute never
    kills an entire message.
    """
    # Subject ──────────────────────────────────────────────────────
    subject = _safe_str(getattr(message, "subject", ""), max_len=500)

    # Sender ───────────────────────────────────────────────────────
    try:
        sender_name  = _safe_str(getattr(message, "sender_name", ""))
        sender_email = _safe_str(getattr(message, "sender_email_address", ""))
        if sender_name and sender_email:
            sender = f"{sender_name} <{sender_email}>"
        else:
            sender = sender_name or sender_email or "Unknown Sender"
    except Exception:
        sender = "Unknown Sender"
    sender = sender[:200]

    # Recipients ───────────────────────────────────────────────────
    to_field = ""
    try:
        count = getattr(message, "recipients_count", 0) or 0
        recipients: list[str] = []
        for i in range(min(count, 100)):
            try:
                r = message.recipients[i]
                addr = _safe_str(getattr(r, "email_address", ""))
                name = _safe_str(getattr(r, "name", ""))
                recipients.append(addr or name)
            except Exception:
                continue
        to_field = "; ".join(filter(None, recipients))[:1000]
    except Exception:
        to_field = ""

    # Date ─────────────────────────────────────────────────────────
    date_str = (
        _format_dt(getattr(message, "delivery_time", None))
        or _format_dt(getattr(message, "client_submit_time", None))
    )

    # Size (bytes) ─────────────────────────────────────────────────
    size = 0
    for attr in ("plain_text_body", "html_body", "transport_headers", "subject"):
        try:
            content = getattr(message, attr, None)
            if content:
                # pypff often returns bytes; len() works on both str and bytes
                size += len(content)
        except Exception:
            continue

    # Attachments count ────────────────────────────────────────────
    attachments = 0
    try:
        attachments = getattr(message, "number_of_attachments", 0) or 0
    except Exception:
        pass

    # Message class (e.g. IPM.Note, IPM.Appointment …) ────────────
    msg_class = _safe_str(getattr(message, "message_class", ""), max_len=100)

    return {
        "Folder":      str(folder_path)[:200],
        "Subject":     subject,
        "From":        sender,
        "To":          to_field,
        "Date":        date_str,
        "Size":        str(size),
        "Attachments": str(attachments),
        "MessageClass": msg_class,
    }


# ────────────────────────── counting ────────────────────────────

def count_total_messages(folder) -> int:
    """Recursively count every message in the folder tree (for the progress bar)."""
    total = 0
    try:
        total += getattr(folder, "number_of_sub_messages", 0) or 0
        for i in range(getattr(folder, "number_of_sub_folders", 0) or 0):
            try:
                sub = folder.get_sub_folder(i)
                if sub:
                    total += count_total_messages(sub)
            except Exception:
                continue
    except Exception:
        pass
    return total


# ───────────────────────── folder walk ──────────────────────────

def process_folder(folder, csv_writer, pbar, counters: dict):
    """
    Recursively walk all folders and write one CSV row per message.
    Uses an iterative approach for subfolders to avoid Python recursion limits
    on deeply nested PST files.
    """
    folder_stack = [folder]

    while folder_stack:
        current = folder_stack.pop()
        try:
            folder_path = get_folder_path(current)
        except Exception:
            folder_path = "Unknown"

        # ── messages in this folder ──
        n_msgs = getattr(current, "number_of_sub_messages", 0) or 0
        for i in range(n_msgs):
            try:
                msg = current.get_sub_message(i)
                if msg is None:
                    continue
                info = extract_email_info(msg, folder_path)
                csv_writer.writerow([
                    info["Folder"],
                    info["Subject"],
                    info["From"],
                    info["To"],
                    info["Date"],
                    info["Size"],
                    info["Attachments"],
                    info["MessageClass"],
                ])
                counters["ok"] += 1
            except Exception:
                counters["err"] += 1
            finally:
                pbar.update(1)

        # ── queue subfolders ──
        n_subs = getattr(current, "number_of_sub_folders", 0) or 0
        for i in range(n_subs):
            try:
                sub = current.get_sub_folder(i)
                if sub:
                    folder_stack.append(sub)
            except Exception:
                continue


# ───────────────────────── PST export ───────────────────────────

FIELDNAMES = ["Folder", "Subject", "From", "To", "Date", "Size", "Attachments", "MessageClass"]


def export_pst_to_csv(pst_file_path: str, output_file_path: str) -> bool:
    """
    Open a single PST file, walk its entire folder tree, and write metadata
    for every message to a UTF-8 CSV file.

    Returns True on success, False on failure.
    """
    if not os.path.exists(pst_file_path):
        print(f"  ❌ PST file not found: {pst_file_path}")
        return False

    # ── open PST ──────────────────────────────────────────────────
    try:
        pst = pypff.file()
        pst.open(pst_file_path)
    except Exception as exc:
        print(f"  ❌ Cannot open PST: {exc}")
        return False

    try:
        root = pst.get_root_folder()
        if root is None:
            print("  ❌ Cannot access root folder")
            return False

        base = os.path.basename(pst_file_path)
        print(f"  📊 Counting messages in {base} …")
        total = count_total_messages(root)
        print(f"  ✓  {total:,} messages found")

        counters = {"ok": 0, "err": 0}

        with open(output_file_path, "w", newline="", encoding="utf-8") as fh:
            writer = csv.writer(fh)
            writer.writerow(FIELDNAMES)

            bar_fmt = "{desc}: {percentage:3.0f}%|{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]"
            with tqdm(total=total, desc=f"  Processing {base}", unit="email",
                      bar_format=bar_fmt) as pbar:
                process_folder(root, writer, pbar, counters)

        print(f"  ✓  Written {counters['ok']:,} rows  |  Errors skipped: {counters['err']:,}")
        return True

    except Exception as exc:
        print(f"  ❌ Critical error: {exc}")
        return False
    finally:
        try:
            pst.close()
        except Exception:
            pass


# ────────────────────── batch processing ────────────────────────

def find_pst_files(directory: str) -> list[str]:
    """Return sorted list of .pst files in *directory* (non-recursive)."""
    return sorted(glob.glob(os.path.join(directory, "*.pst")))


def process_all_pst_files(directory: str) -> None:
    """Discover and process every PST file found in *directory*."""
    sep = "=" * 80

    print(sep)
    print("🔍  SEARCHING FOR PST FILES")
    print(sep)

    pst_files = find_pst_files(directory)
    if not pst_files:
        print(f"❌  No PST files found in: {directory}")
        return

    total_files = len(pst_files)
    print(f"✓   Found {total_files} PST file(s)")
    print()
    print("📁  PST Files:")
    for idx, f in enumerate(pst_files, 1):
        size_mb = os.path.getsize(f) / (1024 ** 2)
        print(f"  {idx:>3}. {os.path.basename(f)}  ({size_mb:.1f} MB)")
    print()

    print(sep)
    print("🚀  PROCESSING")
    print(sep)

    ok_count = 0
    fail_count = 0

    for idx, pst_file in enumerate(pst_files, 1):
        remaining = total_files - idx
        print()
        print("─" * 80)
        print(f"📊  {idx}/{total_files}  |  ✓ Done: {ok_count}  |  ❌ Failed: {fail_count}  |  ⏳ Remaining: {remaining}")
        print(f"🔄  {os.path.basename(pst_file)}")

        base_name   = os.path.splitext(os.path.basename(pst_file))[0]
        output_file = os.path.join(directory, f"{base_name}_email_list.csv")

        if export_pst_to_csv(pst_file, output_file):
            ok_count += 1
            rows = sum(1 for _ in open(output_file, encoding="utf-8")) - 1  # minus header
            print(f"  ✅  Saved → {os.path.basename(output_file)}  ({rows:,} rows)")
        else:
            fail_count += 1
            print("  ❌  Failed")

    # ── summary ───────────────────────────────────────────────────
    print()
    print(sep)
    print("📈  FINAL SUMMARY")
    print(sep)
    print(f"  PST files found:       {total_files}")
    print(f"  ✓  Processed:          {ok_count}")
    print(f"  ❌  Failed:             {fail_count}")
    if total_files:
        print(f"  Success rate:          {ok_count / total_files * 100:.1f}%")
    print(sep)

    if ok_count == total_files:
        print("🎉  All files processed successfully!")
    elif ok_count:
        print("⚠️   Some files failed – check errors above.")
    else:
        print("❌  No files were processed. Check errors above.")
    print(sep)


# ─────────────────────────── entry point ────────────────────────

if __name__ == "__main__":
    if len(sys.argv) > 1:
        directory = sys.argv[1]
    else:
        directory = os.path.dirname(os.path.abspath(__file__))
        print(f"ℹ️   No directory specified – using script location: {directory}")
        print(f"ℹ️   Tip: python {os.path.basename(__file__)} <directory_path>")
        print()

    if not os.path.isdir(directory):
        print(f"❌  Not a valid directory: {directory}")
        sys.exit(1)

    print(f"📂  Working directory: {directory}")
    print("ℹ️   Corrupt or unsupported messages are skipped automatically.")
    print()

    process_all_pst_files(directory)