Refactor ConnectionProxy and tests for improved transient error handling and retry logic

.github/workflows/ci.yml

@@ -153,5 +153,5 @@ jobs:
 
       - name: Check imports
         run: |
-          cd src && python -c "import utils.imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
+          cd src && python -c "import imap_common; import migrate_imap_emails; import backup_imap_emails; import restore_imap_emails; import count_imap_emails; import compare_imap_folders"
           echo "All imports successful"

.github/workflows/publish.yml

@@ -15,14 +15,6 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: "3.x"
-    
-    - name: Update version to match release
-      run: |
-        VERSION=${GITHUB_REF_NAME#v}
-        echo "Updating version to $VERSION"
-        sed -i "s/^version = .*/version = \"$VERSION\"/" pyproject.toml
-        grep "^version =" pyproject.toml
-        
     - name: Install pypa/build
       run: >-
         python3 -m pip install build --user

README.md

@@ -16,7 +16,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### The Scripts
 
-1. **`imap_migrate.py`** (The Solution)
+1. **`migrate_imap_emails.py`** (The Solution)
    - Migrates emails folder-by-folder.
    - **Multi-threaded**: Uses a thread pool to copy messages in parallel for high speed.
   - **Smart De-duplication**: Checks if a message already exists in the destination (matching Message-ID) and skips it if found.
@@ -28,17 +28,17 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Incremental Cache**: Supports `--migrate-cache <path>` to store a local map of processed emails, drastically speeding up subsequent runs by skipping known messages.
    - **Configurable**: Adjustable concurrency and batch sizes to respect server rate limits.
 
-2. **`imap_compare.py`** (The Validator)
+2. **`compare_imap_folders.py`** (The Validator)
    - Connects to both Source and Destination accounts.
    - Prints a side-by-side comparison table of message counts for every folder.
   - Supports comparing IMAP to a local backup folder (`.eml` files) as either the source or destination.
   - Essential for verifying that the migration was successful and that counts match.
 
-3. **`imap_count.py`** (The Investigator)
+3. **`count_imap_emails.py`** (The Investigator)
   - Counts emails in all folders for a single IMAP account (initial assessment / sizing).
   - Also supports counting a local backup folder (`.eml` files) via `--path` (or `BACKUP_LOCAL_PATH`).
 
-4. **`imap_backup.py`** (The Backup)
+4. **`backup_imap_emails.py`** (The Backup)
    - Downloads emails from an IMAP account to a local disk.
    - **Format**: Saves emails as individual `.eml` files (RFC 5322), compatible with Outlook, Thunderbird, and Apple Mail.
    - **Structure**: Replicates the IMAP folder hierarchy locally.
@@ -46,7 +46,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
    - **Sync Mode**: Optionally deletes local `.eml` files that no longer exist on the server (`--dest-delete`).
    - **Gmail Labels Preservation**: Creates a `labels_manifest.json` file mapping each email's Message-ID to its Gmail labels, enabling proper restoration with labels intact.
 
-5. **`imap_restore.py`** (The Restore)
+5. **`restore_imap_emails.py`** (The Restore)
    - Uploads emails from a local backup to an IMAP server.
    - **Format**: Reads `.eml` files and uploads them preserving original dates.
    - **Structure**: Recreates the folder hierarchy on the destination server.
@@ -67,36 +67,7 @@ This repository contains a set of Python scripts designed to migrate emails betw
 
 ### 2. Installation
 
-#### Installation via PyPI (Recommended)
-You can install the tools directly from PyPI.
-
-**For macOS Users (and other externally managed environments):**
-Due to recent changes in Python environments on macOS, it is recommended to use `pipx` to install the tools globally without conflicting with system packages.
-
-1. Install `pipx` (if not already installed):
-   ```bash
-   brew install pipx
-   pipx ensurepath
-   ```
-2. Install the tools:
-   ```bash
-   pipx install imap-migration-tools
-   ```
-
-**For Standard Environments:**
-```bash
-pip install imap-migration-tools
-```
-
-Once installed via PyPI, the following commands are globally available in your terminal:
-- `imap-backup`
-- `imap-compare`
-- `imap-count`
-- `imap-migrate`
-- `imap-restore`
-
-#### Installation from Source
-Alternatively, clone this repository or download the script files to your local machine.
+Clone this repository or download the script files to your local machine.
 
 #### macOS
 macOS often comes with Python, but it's best to install the latest version.
@@ -158,7 +129,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```bash
-   python3 imap_migrate.py
+   python3 migrate_imap_emails.py
    ```
 
 #### Windows (PowerShell)
@@ -175,7 +146,7 @@ You can configure the scripts using **Environment Variables** (recommended for s
 
 2. **Run:**
    ```powershell
-   python imap_migrate.py
+   python migrate_imap_emails.py
    ```
 
 ### Method 2: Command Line Arguments (Overrides)
@@ -183,7 +154,7 @@ All scripts support command-line arguments which take precedence over environmen
 
 **Migration:**
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -196,7 +167,7 @@ python3 imap_migrate.py \
 
 **Comparison:**
 ```bash
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -205,14 +176,14 @@ python3 imap_compare.py \
   --dest-pass "your-app-password"
 
 # Compare IMAP source to a local backup folder
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./my_backup"
 
 # Compare a local backup folder to an IMAP destination
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -221,23 +192,23 @@ python3 imap_compare.py \
 
 **Counting:**
 ```bash
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --pass "secret"
 
 # Count a local backup folder
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --path "./my_backup"
 
 # Or via environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 imap_count.py
+python3 count_imap_emails.py
 ```
 
 **Counting (OAuth2):**
 ```bash
-python3 imap_count.py \
+python3 count_imap_emails.py \
    --host "imap.gmail.com" \
    --user "me@gmail.com" \
    --oauth2-client-id "id" \
@@ -248,12 +219,12 @@ export IMAP_HOST="imap.gmail.com"
 export IMAP_USERNAME="me@gmail.com"
 export OAUTH2_CLIENT_ID="id"
 export OAUTH2_CLIENT_SECRET="secret"  # Required for Google
-python3 imap_count.py
+python3 count_imap_emails.py
 ```
 
 **Backup:**
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "me@gmail.com" \
   --src-pass "your-app-password" \
@@ -265,7 +236,7 @@ python3 imap_backup.py \
 ### 1. Full Migration
 Migrate all folders from Source to Destination.
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -278,7 +249,7 @@ python3 imap_migrate.py \
 For Gmail -> Gmail migrations, `--gmail-mode` migrates only `[Gmail]/All Mail` (no duplicates) and applies labels by copying messages into label folders.
 
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --gmail-mode \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
@@ -292,7 +263,7 @@ python3 imap_migrate.py \
 Use a local cache file to remember processed emails. Use this for large migrations that may be interrupted or need to run multiple times.
 
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.source.com" \
   --src-user "source" \
   --src-pass "pass" \
@@ -310,7 +281,7 @@ Preserve IMAP flags (`\Seen`, `\Flagged`, `\Answered`, `\Draft`) during migratio
 If an email already exists on the destination (duplicate), the script can still sync missing flags on the existing message.
 
 ```bash
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --preserve-flags \
   --src-host "imap.example.com" \
   --src-user "source@example.com" \
@@ -323,8 +294,8 @@ python3 imap_migrate.py \
 ### 2. Single Folder Migration
 Migrate ONLY a specific folder (e.g., trying to fix just "Important" or "Sent").
 ```bash
-# Syntax: python3 imap_migrate.py "[Folder Name]"
-python3 imap_migrate.py \
+# Syntax: python3 migrate_imap_emails.py "[Folder Name]"
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -338,7 +309,7 @@ python3 imap_migrate.py \
 Migrate and **delete** from source immediately after verifying the copy.
 ```bash
 # Using flag
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -348,7 +319,7 @@ python3 imap_migrate.py \
   --src-delete
 
 # Or specific folder with delete
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -365,7 +336,7 @@ Keep destination in sync by deleting emails that no longer exist in the source.
 Note: `--dest-delete` is not supported in `--gmail-mode`.
 ```bash
 # Migration: Delete destination emails not found in source
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -375,7 +346,7 @@ python3 imap_migrate.py \
   --dest-delete
 
 # Backup: Delete local .eml files not found on server
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -383,7 +354,7 @@ python3 imap_backup.py \
   --dest-delete
 
 # Restore: Delete server emails not found in local backup
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -396,7 +367,7 @@ python3 imap_restore.py \
 ### 5. Verify Migration
 Compare counts between source and destination.
 ```bash
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
@@ -405,14 +376,14 @@ python3 imap_compare.py \
   --dest-pass "dest-app-password"
 
 # IMAP source -> local backup destination
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "source@gmail.com" \
   --src-pass "source-app-password" \
   --dest-path "./my_backup"
 
 # local backup source -> IMAP destination
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "dest@domain.com" \
@@ -430,21 +401,21 @@ INBOX                   | 1250         | 1250       | MATCH
 Download all your emails to your computer as `.eml` files.
 ```bash
 # Backup all folders from an IMAP account to a local folder
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./backup_folder"
 
 # Or via command line
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "/Users/jdoe/Documents/Emails"
 
 # Backup single folder
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -453,18 +424,18 @@ python3 imap_backup.py \
 ```
 
 ### 6a. Compare IMAP vs Local Backup
-Use `imap_compare.py` to validate an IMAP account against a local backup created by `imap_backup.py`.
+Use `compare_imap_folders.py` to validate an IMAP account against a local backup created by `backup_imap_emails.py`.
 
 ```bash
 # Option 1: IMAP source -> local destination
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
   --dest-path "./my_backup"
 
 # Option 2: local source -> IMAP destination
-python3 imap_compare.py \
+python3 compare_imap_folders.py \
   --src-path "./my_backup" \
   --dest-host "imap.other.com" \
   --dest-user "you@domain.com" \
@@ -479,15 +450,15 @@ export DEST_LOCAL_PATH="./my_backup"
 ```
 
 ### 6b. Count a Local Backup
-Use `imap_count.py` to get per-folder counts from a local backup created by `imap_backup.py`.
+Use `count_imap_emails.py` to get per-folder counts from a local backup created by `backup_imap_emails.py`.
 
 ```bash
 # Option 1: explicit path
-python3 imap_count.py --path "./my_backup"
+python3 count_imap_emails.py --path "./my_backup"
 
 # Option 2: environment variable
 export BACKUP_LOCAL_PATH="./my_backup"
-python3 imap_count.py
+python3 count_imap_emails.py
 ```
 
 ### 7. Gmail Backup with Labels Preservation
@@ -495,7 +466,7 @@ When backing up a Gmail account, use `--gmail-mode` for the recommended workflow
 
 ```bash
 # Recommended: Use --gmail-mode for simplest workflow
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -505,7 +476,7 @@ python3 imap_backup.py \
 
 This is equivalent to the more verbose:
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -517,7 +488,7 @@ python3 imap_backup.py \
 **For large accounts (100K+ emails)**, you can build the manifest first to test:
 ```bash
 # Step 1: Build manifest only (fast, no download)
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -525,7 +496,7 @@ python3 imap_backup.py \
   --manifest-only
 
 # Step 2: Download emails (can run later, manifest already exists)
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-pass "your-app-password" \
@@ -571,7 +542,7 @@ python3 imap_backup.py \
 For non-Gmail servers, you can preserve read/starred status with `--preserve-flags`:
 
 ```bash
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.example.com" \
   --src-user "you@example.com" \
   --src-pass "your-password" \
@@ -590,14 +561,14 @@ Use `--full-restore` to force the legacy behavior (process all emails and re-syn
 
 ```bash
 # Restore all folders from backup
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
   --dest-pass "your-app-password"
 
 # Force full restore (legacy behavior)
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -605,7 +576,7 @@ python3 imap_restore.py \
   --full-restore
 
 # Restore with flags (read/starred status)
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.example.com" \
   --dest-user "you@example.com" \
@@ -613,7 +584,7 @@ python3 imap_restore.py \
   --apply-flags
 
 # Restore a specific folder
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./my_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "you@gmail.com" \
@@ -625,7 +596,7 @@ python3 imap_restore.py \
 Restore a Gmail backup with full label structure using `--gmail-mode`:
 
 ```bash
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -642,7 +613,7 @@ python3 imap_restore.py \
 
 **Alternatively**, restore folders individually with labels and flags applied:
 ```bash
-python3 imap_restore.py \
+python3 restore_imap_emails.py \
   --src-path "./gmail_backup" \
   --dest-host "imap.gmail.com" \
   --dest-user "newaccount@gmail.com" \
@@ -664,7 +635,7 @@ To use OAuth2, pass `--oauth2-client-id` (or `--src-oauth2-client-id`/`--dest-oa
 
 Environment variable equivalents:
 - Dual-account scripts: `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET` and `DEST_OAUTH2_CLIENT_ID`, `DEST_OAUTH2_CLIENT_SECRET`.
-- Single-account scripts (like `imap_count.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
+- Single-account scripts (like `count_imap_emails.py`): `OAUTH2_CLIENT_ID`, `OAUTH2_CLIENT_SECRET` (also accepts `SRC_OAUTH2_CLIENT_ID`, `SRC_OAUTH2_CLIENT_SECRET`).
 
 ### Microsoft (Outlook / Office 365)
 
@@ -703,7 +674,7 @@ Requires the `msal` package (`pip install msal`). Uses the **device code flow**
 pip install msal
 
 # Migration with Microsoft OAuth2 on source
-python3 imap_migrate.py \
+python3 migrate_imap_emails.py \
   --src-host "outlook.office365.com" \
   --src-user "user@contoso.com" \
   --src-oauth2-client-id "your-azure-app-client-id" \
@@ -723,7 +694,7 @@ Requires the `google-auth-oauthlib` package (`pip install google-auth-oauthlib`)
 pip install google-auth-oauthlib
 
 # Backup with Google OAuth2
-python3 imap_backup.py \
+python3 backup_imap_emails.py \
   --src-host "imap.gmail.com" \
   --src-user "you@gmail.com" \
   --src-oauth2-client-id "your-google-client-id" \
@@ -736,7 +707,7 @@ The script will open your default browser for Google sign-in. After authorizing,
 ## Troubleshooting
 
 - **"Too many simultaneous connections"**:
-  IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `imap_migrate.py` uses multiple threads, you may hit this limit.
+  IMAP servers (especially Gmail) limit the number of active connections per IP or user (typically ~15). Since `migrate_imap_emails.py` uses multiple threads, you may hit this limit.
   **Solution**: Reduce `MAX_WORKERS` to `4` or `2` using the environment variable.
 
 - **Authentication Errors**:
@@ -789,7 +760,7 @@ PYTHONPATH=src pytest test/ -v
 make coverage
 
 # Run a specific test file
-PYTHONPATH=src pytest test/test_imap_migrate.py -v
+PYTHONPATH=src pytest test/test_migrate_imap_emails.py -v
 
 # Run a specific test
 PYTHONPATH=src pytest test/test_imap_common.py::TestNormalizeFolderName -v
@@ -821,11 +792,11 @@ make ci
 
 | Test File | Description |
 |-----------|-------------|
-| `test_imap_migrate.py` | Email migration tests (basic, duplicates, deletion, folders) |
-| `test_imap_backup.py` | Backup functionality tests |
-| `test_imap_restore.py` | Restore functionality tests |
-| `test_imap_count.py` | Email counting tests |
-| `test_imap_compare.py` | Folder comparison tests |
+| `test_migrate_imap_emails.py` | Email migration tests (basic, duplicates, deletion, folders) |
+| `test_backup_imap_emails.py` | Backup functionality tests |
+| `test_restore_imap_emails.py` | Restore functionality tests |
+| `test_count_imap_emails.py` | Email counting tests |
+| `test_compare_imap_folders.py` | Folder comparison tests |
 | `test_imap_common.py` | Shared utility function tests |
 
 ### Continuous Integration

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imap-migration-tools"
-version = "1.1.0"
+version = "1.0.0"
 authors = [
   { name="Javier Callico", email="jcallico@callicode.com" },
   { name="Nathan Moinvaziri", email="nathan@nathanm.com" }
@@ -30,11 +30,11 @@ dependencies = [
 "Bug Tracker" = "https://github.com/jcallico/imap-migration-tools/issues"
 
 [project.scripts]
-imap-backup = "imap_backup:main"
-imap-restore = "imap_restore:main"
-imap-migrate = "imap_migrate:main"
-imap-count = "imap_count:main"
-imap-compare = "imap_compare:main"
+imap-backup = "backup_imap_emails:main"
+imap-restore = "restore_imap_emails:main"
+imap-migrate = "migrate_imap_emails:main"
+imap-count = "count_imap_emails:main"
+imap-compare = "compare_imap_folders:main"
 
 [tool.setuptools.package-dir]
 "" = "src"

src/backup_imap_emails.py

@@ -1,24 +1,895 @@
-#!/usr/bin/env python3
 """
-DEPRECATED: This script has been renamed.
-Please use imap_backup.py instead.
+IMAP Email Backup Script
+
+Backs up emails from an IMAP account to a local directory.
+Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
+most email clients (Thunderbird, Apple Mail, Outlook, etc.).
+
+Features:
+- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
+- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
+- Folder Replication: Recreates the IMAP folder structure locally.
+- Parallel Processing: Uses multithreading for fast downloads.
+- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
+
+Configuration (Environment Variables):
+    SRC_IMAP_HOST, SRC_IMAP_USERNAME: Source credentials.
+    SRC_IMAP_PASSWORD: Source password (or App Password).
+
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID: OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
+
+  BACKUP_LOCAL_PATH: Destination local directory.
+  MAX_WORKERS: Number of concurrent threads (default: 10).
+  BATCH_SIZE: Number of emails to process per batch (default: 10).
+  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
+  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
+  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
+  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
+  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
+              Default is "false".
+
+Usage:
+    python3 backup_imap_emails.py \
+        --src-host "imap.example.com" \
+        --src-user "you@example.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup"
+
+Gmail Labels:
+    python3 backup_imap_emails.py \
+        --src-host "imap.gmail.com" \
+        --src-user "you@gmail.com" \
+        --src-pass "your-app-password" \
+        --dest-path "./my_backup" \
+        --preserve-labels \
+        "[Gmail]/All Mail"
+  This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
+  file that maps each email's Message-ID to its Gmail labels for later restoration.
 """
 
+import argparse
+import concurrent.futures
+import json
+import os
 import sys
+import threading
 
-from imap_backup import main
+import imap_common
+import imap_oauth2
+import imap_session
+import provider_exchange
+import provider_gmail
+
+# Defaults
+MAX_WORKERS = 10
+BATCH_SIZE = 10
+MANIFEST_FILENAME = "labels_manifest.json"
+
+# Thread-local storage
+thread_local = threading.local()
+safe_print = imap_common.safe_print
+
+
+def process_single_uid(src, uid, folder_name, local_folder_path):
+    """
+    Fetch and save a single email by UID.
+
+    Args:
+        src: IMAP connection
+        uid: Message UID to fetch
+        folder_name: Current folder name (for logging)
+        local_folder_path: Local directory to save email
+
+    Returns:
+        Tuple of (success: bool, connection):
+        - (True, src): UID processed successfully (or skipped)
+        - (False, src): Auth error occurred, caller should retry after reconnect
+    """
+    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
 
-if __name__ == "__main__":
-    print(
-        "WARNING: 'backup_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_backup.py'.",
-        file=sys.stderr,
-    )
-    print("Redirecting to 'imap_backup.py'...", file=sys.stderr)
     try:
-        main()
+        resp, data = src.uid("fetch", uid, "(RFC822)")
+        if resp != "OK" or not data or data[0] is None:
+            safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
+            return (True, src)  # Don't retry, move on
+
+        raw_email = None
+        for item in data:
+            if isinstance(item, tuple):
+                raw_email = item[1]
+                break
+
+        # Derive Subject for filename from the already-fetched message bytes.
+        _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
+        if not subject:
+            clean_subject = "No Subject"
+        else:
+            clean_subject = imap_common.sanitize_filename(subject)
+            clean_subject = clean_subject[:100]
+
+        filename = f"{uid_str}_{clean_subject}.eml"
+        full_path = os.path.join(local_folder_path, filename)
+
+        if os.path.exists(full_path):
+            return (True, src)  # Already exists
+
+        if raw_email:
+            try:
+                with open(full_path, "wb") as f:
+                    f.write(raw_email)
+                safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
+            except OSError as e:
+                safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
+        else:
+            safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
+
+        return (True, src)
+
+    except Exception as e:
+        if imap_oauth2.is_auth_error(e):
+            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
+            return (False, src)  # Signal retry needed
+        else:
+            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
+            return (True, src)  # Don't retry other errors
+
+
+def process_batch(uids, folder_name, src_conf, local_folder_path):
+    src = imap_session.get_thread_connection(thread_local, "src", src_conf)
+    if not src:
+        safe_print("Error: Could not establish connection for batch.")
+        return
+
+    try:
+        src.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
+        return
+
+    for uid in uids:
+        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
+        max_retries = 2
+
+        for attempt in range(max_retries):
+            src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
+            thread_local.src = src
+            if not ok:
+                safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
+                return
+
+            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
+            thread_local.src = src
+
+            if success:
+                break
+            if attempt < max_retries - 1:
+                src = None
+                thread_local.src = None
+
+
+def get_existing_uids(local_path):
+    """
+    Scans the local directory for files matching pattern matches {UID}_*.eml
+    Returns a set of UIDs (as strings).
+    """
+    existing = set()
+    if not os.path.exists(local_path):
+        return existing
+
+    try:
+        for filename in os.listdir(local_path):
+            if filename.endswith(".eml") and "_" in filename:
+                # Expecting UID_Subject.eml
+                parts = filename.split("_", 1)
+                if parts[0].isdigit():
+                    existing.add(parts[0])
+    except Exception:
+        pass
+    return existing
+
+
+# Standard IMAP flags that can be preserved during migration
+# \Recent is session-specific and cannot be set by clients
+# \Deleted should not be preserved as it marks messages for removal
+PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
+
+
+def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
+    """
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
+    with OAuth2 session management.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_info, imap_conn) where message_info is:
+        { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+        The returned imap_conn may be different if reconnection occurred.
+    """
+    message_info = {}
+
+    # Ensure connection is healthy (refresh OAuth2 token if needed) before initial select
+    if src_conf:
+        imap_conn = imap_session.ensure_connection(imap_conn, src_conf)
+        if not imap_conn:
+            safe_print(f"Could not establish connection for folder {folder_name}")
+            return (message_info, None)
+
+    try:
+        imap_conn.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Could not select folder {folder_name}: {e}")
+        return (message_info, imap_conn)
+
+    try:
+        resp, data = imap_conn.uid("search", None, "ALL")
+        if resp != "OK" or not data or not data[0]:
+            return (message_info, imap_conn)
+
+        uids = data[0].split()
+        if not uids:
+            return (message_info, imap_conn)
+
+        total_uids = len(uids)
+
+        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
+        batch_size = 200
+        for i in range(0, len(uids), batch_size):
+            batch = uids[i : i + batch_size]
+            uid_range = b",".join(batch)
+
+            # Report progress
+            if progress_callback:
+                progress_callback(min(i + batch_size, total_uids), total_uids)
+
+            # Proactively refresh token and ensure folder is selected
+            if src_conf:
+                imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
+                if not ok:
+                    safe_print(f"ERROR: Connection/folder lost in {folder_name}")
+                    break
+
+            try:
+                # Fetch both Message-ID header and FLAGS
+                resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
+                if resp != "OK":
+                    continue
+
+                # Parse response - items come in pairs for each message
+                for item in items:
+                    if isinstance(item, tuple) and len(item) >= 2:
+                        # First element contains UID and FLAGS info
+                        meta_str = (
+                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
+                        )
+
+                        # Extract all preservable flags from the metadata
+                        flags = []
+                        for flag in imap_common.PRESERVABLE_FLAGS:
+                            if flag in meta_str:
+                                flags.append(flag)
+
+                        # Second element contains the header
+                        msg_id = imap_common.extract_message_id(item[1])
+                        if msg_id:
+                            message_info[msg_id] = {"flags": flags}
+            except Exception as e:
+                safe_print(f"Error fetching batch in {folder_name}: {e}")
+                continue
+
+    except Exception as e:
+        safe_print(f"Error searching folder {folder_name}: {e}")
+
+    return (message_info, imap_conn)
+
+
+def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
+    """
+    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
+
+    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
+    return message_info
+
+
+def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
+    """
+    Returns a set of Message-IDs for all emails in a given folder,
+    with OAuth2 session management.
+
+    Args:
+        imap_conn: IMAP connection
+        folder_name: Folder to scan
+        src_conf: Connection config dict for OAuth2 refresh
+        progress_callback: Optional callback(current, total) for progress reporting
+
+    Returns:
+        Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
+        The returned imap_conn may be different if reconnection occurred.
+    """
+    info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
+    return (set(info.keys()), imap_conn)
+
+
+def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
+    """
+    Returns a set of Message-IDs for all emails in a given folder.
+    This is a convenience wrapper around get_message_info_in_folder.
+    Optional progress_callback(current, total) for progress reporting.
+    """
+    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
+    return set(info.keys())
+
+
+def build_labels_manifest(imap_conn, local_path, src_conf=None):
+    """
+    Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
+    Scans all folders (labels) in the account and records which Message-IDs
+    appear in each label, plus their flags from [Gmail]/All Mail.
+
+    Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
+    Saves the manifest to labels_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
+    """
+    import time
+
+    manifest = {}
+    start_time = time.time()
+    total_emails_scanned = 0
+
+    safe_print("--- Building Gmail Labels Manifest ---")
+
+    # Get all selectable folders and filter to label folders
+    all_folders = imap_common.list_selectable_folders(imap_conn)
+    if not all_folders:
+        safe_print("Error: Could not list folders for label mapping.")
+        return manifest
+
+    # First, scan [Gmail]/All Mail to get the authoritative flags for all emails
+    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
+    all_mail_start = time.time()
+
+    def all_mail_progress_cb(current, total):
+        elapsed = time.time() - start_time
+        print(
+            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+            end="",
+            flush=True,
+        )
+
+    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
+        imap_conn, provider_gmail.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
+    )
+    print()  # New line after progress
+
+    # Initialize manifest with flags from All Mail
+    for msg_id, info in all_mail_info.items():
+        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
+
+    all_mail_elapsed = time.time() - all_mail_start
+    # Count flag statistics
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print(f"  -> {len(all_mail_info)} emails scanned ({all_mail_elapsed:.1f}s)")
+    safe_print(f"  -> Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}\n")
+
+    # Keep connection alive
+    try:
+        imap_conn.noop()
+    except Exception:
+        pass
+
+    # Parse folder names and filter to label folders
+    label_folders = [f for f in all_folders if provider_gmail.is_label_folder(f)]
+
+    total_folders = len(label_folders)
+    safe_print(f"Found {total_folders} label folders to scan.\n")
+
+    # Scan each label folder
+    for folder_idx, folder_name in enumerate(label_folders, 1):
+        folder_start = time.time()
+
+        # Progress callback for this folder
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
+        print()  # New line after progress
+
+        # Keep connection alive between folders
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+        # Determine the label name to store
+        # For [Gmail]/Sent Mail -> "Sent Mail"
+        # For [Gmail]/Starred -> "Starred"
+        # For INBOX -> "INBOX"
+        # For user folders -> folder name as-is
+        if folder_name.startswith("[Gmail]/"):
+            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
+        else:
+            label_name = folder_name
+
+        for msg_id in message_ids:
+            if msg_id not in manifest:
+                # Email not in All Mail (rare, but handle it)
+                manifest[msg_id] = {"labels": [], "flags": []}
+            if label_name not in manifest[msg_id]["labels"]:
+                manifest[msg_id]["labels"].append(label_name)
+
+        folder_elapsed = time.time() - folder_start
+        total_emails_scanned += len(message_ids)
+        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print("\nManifest building complete:")
+    safe_print(f"  - Folders scanned: {total_folders + 1}")  # +1 for All Mail
+    safe_print(f"  - Total email-label mappings: {total_emails_scanned}")
+    safe_print(f"  - Unique emails: {len(manifest)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
+    try:
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2, ensure_ascii=False)
+        safe_print(f"\nLabels manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
+    """
+    Builds a manifest mapping Message-IDs to their IMAP flags.
+    For non-Gmail servers, scans specified folders (or all folders if not specified).
+
+    Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
+    Saves the manifest to flags_manifest.json in the backup directory.
+    Optional src_conf for automatic token refresh on expiration.
+    """
+    import time
+
+    manifest = {}
+    start_time = time.time()
+
+    safe_print("--- Building Flags Manifest ---")
+
+    # Get folders to scan
+    if folders_to_scan:
+        all_folders = folders_to_scan
+    else:
+        try:
+            typ, folders = imap_conn.list()
+            if typ != "OK":
+                safe_print("Error: Could not list folders.")
+                return manifest
+            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
+        except Exception as e:
+            safe_print(f"Error listing folders: {e}")
+            return manifest
+
+    total_folders = len(all_folders)
+    safe_print(f"Found {total_folders} folders to scan.\n")
+
+    # Scan each folder
+    for folder_idx, folder_name in enumerate(all_folders, 1):
+        folder_start = time.time()
+
+        def progress_cb(current, total):
+            elapsed = time.time() - start_time
+            print(
+                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
+                end="",
+                flush=True,
+            )
+
+        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
+        folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
+        print()  # New line after progress
+
+        # Merge into manifest (keep first occurrence of flags)
+        for msg_id, info in folder_info.items():
+            if msg_id not in manifest:
+                manifest[msg_id] = {"flags": info.get("flags", [])}
+
+        folder_elapsed = time.time() - folder_start
+        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
+
+        # Keep connection alive
+        try:
+            imap_conn.noop()
+        except Exception:
+            pass
+
+    # Summary
+    total_elapsed = time.time() - start_time
+    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
+    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
+    safe_print("\nFlags manifest building complete:")
+    safe_print(f"  - Folders scanned: {total_folders}")
+    safe_print(f"  - Unique emails: {len(manifest)}")
+    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
+    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
+
+    # Save manifest
+    manifest_path = os.path.join(local_path, "flags_manifest.json")
+    try:
+        with open(manifest_path, "w", encoding="utf-8") as f:
+            json.dump(manifest, f, indent=2, ensure_ascii=False)
+        safe_print(f"\nFlags manifest saved to: {manifest_path}")
+    except Exception as e:
+        safe_print(f"Error saving manifest: {e}")
+
+    return manifest
+
+
+def delete_orphan_local_files(local_folder_path, server_uids):
+    """
+    Delete local .eml files that no longer exist on the server.
+    Args:
+        local_folder_path: Path to local folder containing .eml files
+        server_uids: Set of UID strings currently on server
+    Returns:
+        Count of deleted files
+    """
+    deleted_count = 0
+    if not os.path.exists(local_folder_path):
+        return deleted_count
+
+    try:
+        for filename in os.listdir(local_folder_path):
+            if not filename.endswith(".eml") or "_" not in filename:
+                continue
+
+            # Extract UID from filename (format: {UID}_{Subject}.eml)
+            parts = filename.split("_", 1)
+            if not parts[0].isdigit():
+                continue
+
+            local_uid = parts[0]
+            if local_uid not in server_uids:
+                file_path = os.path.join(local_folder_path, filename)
+                try:
+                    os.remove(file_path)
+                    safe_print(f"  -> Deleted orphan: {filename}")
+                    deleted_count += 1
+                except Exception as e:
+                    safe_print(f"  -> Error deleting {filename}: {e}")
+    except Exception as e:
+        safe_print(f"Error scanning for orphan files: {e}")
+
+    return deleted_count
+
+
+def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
+    safe_print(f"--- Processing Folder: {folder_name} ---")
+
+    # create local path
+    # Handle folder separators. IMAP output might be "Parent/Child"
+    # We rely on OS to handle "Parent/Child" as subdirectories using join
+    # But clean the segments
+    cleaned_name = folder_name.replace("/", os.sep)
+    local_folder_path = os.path.join(local_base_path, cleaned_name)
+
+    try:
+        os.makedirs(local_folder_path, exist_ok=True)
+    except Exception as e:
+        safe_print(f"Error creating directory {local_folder_path}: {e}")
+        return
+
+    # Select IMAP folder
+    try:
+        src_main.select(f'"{folder_name}"', readonly=True)
+    except Exception as e:
+        safe_print(f"Skipping {folder_name}: {e}")
+        return
+
+    # Search all
+    resp, data = src_main.uid("search", None, "ALL")
+    if resp != "OK":
+        return
+
+    uids = data[0].split()
+    total_on_server = len(uids)
+
+    # Build set of server UIDs for comparison
+    server_uid_set = set()
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        server_uid_set.add(u_str)
+
+    if total_on_server == 0:
+        safe_print(f"Folder {folder_name} is empty.")
+        # If dest_delete enabled, delete all local files
+        if dest_delete:
+            deleted = delete_orphan_local_files(local_folder_path, set())
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+        return
+
+    # Incremental Optimization
+    # Read local directory to find UIDs we already have
+    existing_uids = get_existing_uids(local_folder_path)
+
+    # Delete orphan local files if dest_delete is enabled
+    if dest_delete:
+        orphan_uids = existing_uids - server_uid_set
+        if orphan_uids:
+            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
+            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
+            if deleted > 0:
+                safe_print(f"Deleted {deleted} orphan files from local backup.")
+
+    # Filter UIDs
+    # decode uid first if bytes
+    uids_to_download = []
+    for u in uids:
+        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
+        if u_str not in existing_uids:
+            uids_to_download.append(u)
+
+    skipped = total_on_server - len(uids_to_download)
+    if skipped > 0:
+        safe_print(f"Skipping {skipped} emails (already exist locally).")
+
+    if not uids_to_download:
+        safe_print(f"Folder {folder_name} is up to date.")
+        return
+
+    safe_print(f"Downloading {len(uids_to_download)} new emails...")
+
+    # Create batches
+    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
+
+    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
+    try:
+        futures = []
+        for batch in uid_batches:
+            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
+
+        for future in concurrent.futures.as_completed(futures):
+            future.result()
+
     except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
+        executor.shutdown(wait=False, cancel_futures=True)
+        raise
+    finally:
+        executor.shutdown(wait=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
+
+    # Source
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+
+    # Authentication: require either password OR OAuth2 client-id (unless provided via env vars)
+    auth_required = not bool(default_src_pass or default_src_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    # OAuth2
+    auth_group.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
+
+    # Destination (Local Path)
+    env_path = os.getenv("BACKUP_LOCAL_PATH")
+    parser.add_argument(
+        "--dest-path",
+        default=env_path,
+        required=not bool(env_path),
+        help="Local destination path (or BACKUP_LOCAL_PATH)",
+    )
+
+    # Config
+    parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
+    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
+
+    # Gmail Labels
+    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-labels",
+        action="store_true",
+        default=env_preserve_labels,
+        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
+    )
+    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
+    parser.add_argument(
+        "--preserve-flags",
+        action="store_true",
+        default=env_preserve_flags,
+        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
+    )
+    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
+    parser.add_argument(
+        "--manifest-only",
+        action="store_true",
+        default=env_manifest_only,
+        help="Gmail only: Build the labels manifest and exit without downloading emails",
+    )
+    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
+    parser.add_argument(
+        "--gmail-mode",
+        action="store_true",
+        default=env_gmail_mode,
+        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
+    )
+
+    # Sync mode: delete local files not on server
+    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
+    parser.add_argument(
+        "--dest-delete",
+        action="store_true",
+        default=env_dest_delete,
+        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
+    )
+
+    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
+
+    args = parser.parse_args()
+
+    global MAX_WORKERS, BATCH_SIZE
+    MAX_WORKERS = args.workers
+    BATCH_SIZE = args.batch
+
+    # Build connection config (acquires OAuth2 token if configured)
+    src_conf = imap_session.build_imap_conf(
+        args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
+    )
+
+    # Expand path (~/...)
+    local_path = os.path.expanduser(args.dest_path)
+    if not os.path.exists(local_path):
+        try:
+            os.makedirs(local_path)
+            print(f"Created backup directory: {local_path}")
+        except Exception as e:
+            print(f"Error creating backup directory: {e}")
+            sys.exit(1)
+
+    print("\n--- Configuration Summary ---")
+    print(f"Source Host     : {args.src_host}")
+    print(f"Source User     : {args.src_user}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
+    print(f"Destination Path: {local_path}")
+    if args.gmail_mode:
+        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
+    elif args.manifest_only:
+        print("Mode            : Manifest Only (no email download)")
+    elif args.folder:
+        print(f"Target Folder   : {args.folder}")
+    if args.preserve_labels or args.manifest_only or args.gmail_mode:
+        print("Preserve Labels : Yes (Gmail)")
+    if args.preserve_flags or args.gmail_mode:
+        print("Preserve Flags  : Yes (read/starred/answered/draft)")
+    if args.dest_delete:
+        print("Dest Delete     : Yes (remove local orphans)")
+    print("-----------------------------\n")
+
+    try:
+        src = imap_common.get_imap_connection_from_conf(src_conf)
+        if not src:
+            sys.exit(1)
+
+        # Build labels manifest BEFORE backing up emails (Gmail mode)
+        # This way we capture the label state at backup time
+        if args.preserve_labels or args.manifest_only or args.gmail_mode:
+            print("Building Gmail labels manifest...")
+            print("This scans all folders to map Message-IDs to labels and flags.\n")
+            build_labels_manifest(src, local_path, src_conf)
+            print("")  # Blank line after manifest building
+        # Build flags-only manifest for non-Gmail servers
+        elif args.preserve_flags:
+            print("Building flags manifest...")
+            print("This scans folders to capture read/starred/etc status.\n")
+            # Get folders to scan
+            folders_to_scan = [args.folder] if args.folder else None
+            build_flags_manifest(src, local_path, folders_to_scan, src_conf)
+            print("")  # Blank line after manifest building
+
+        # If manifest-only mode, we're done
+        if args.manifest_only:
+            try:
+                src.logout()
+            except Exception:
+                pass  # Connection may already be closed
+            manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
+            print("\nManifest-only mode complete.")
+            print(f"Labels manifest saved to: {manifest_path}")
+            print("\nTo download emails, run again without --manifest-only:")
+            print(f'  python3 backup_imap_emails.py --dest-path "{local_path}" "[Gmail]/All Mail"')
+            sys.exit(0)
+
+        # Gmail mode: backup only [Gmail]/All Mail
+        if args.gmail_mode:
+            backup_folder(src, provider_gmail.GMAIL_ALL_MAIL, local_path, src_conf, args.dest_delete)
+        elif args.folder:
+            backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
+        else:
+            # Reconnect after potentially long manifest building
+            src = imap_session.ensure_connection(src, src_conf)
+            if not src:
+                print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
+                sys.exit(0)
+            folders = imap_common.list_selectable_folders(src)
+            for name in folders:
+                if provider_exchange.is_special_folder(name):
+                    print(f"Skipping Exchange system folder: {name}")
+                    continue
+                src = imap_session.ensure_connection(src, src_conf)
+                if not src:
+                    print("Fatal: Could not reconnect to IMAP server. Aborting.")
+                    sys.exit(1)
+                backup_folder(src, name, local_path, src_conf, args.dest_delete)
+
+        try:
+            src.logout()
+        except Exception:
+            pass  # Connection may already be closed
+        print("\nBackup completed successfully.")
+
+        if args.preserve_labels or args.gmail_mode:
+            manifest_path = os.path.join(local_path, "labels_manifest.json")
+            print(f"\nGmail labels manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply labels and flags to emails.")
+        elif args.preserve_flags:
+            manifest_path = os.path.join(local_path, "flags_manifest.json")
+            print(f"\nFlags manifest saved to: {manifest_path}")
+            print("Use this file when restoring to reapply read/starred status to emails.")
+
+    except KeyboardInterrupt:
+        print("\nBackup interrupted by user.")
         sys.exit(0)
     except Exception as e:
         print(f"Fatal Error: {e}")
-        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

src/compare_imap_folders.py

@@ -1,24 +1,329 @@
-#!/usr/bin/env python3
 """
-DEPRECATED: This script has been renamed.
-Please use imap_compare.py instead.
+IMAP Folder Comparison Script
+
+This script compares email counts between a source and a destination.
+Each side can be either an IMAP account or a local backup folder.
+It iterates through all folders found in the source account and checks the corresponding
+folder in the destination account.
+
+Configuration (Environment Variables):
+  Source Account:
+    SRC_IMAP_HOST       : Source IMAP Host
+    SRC_IMAP_USERNAME   : Source Username/Email
+    SRC_IMAP_PASSWORD   : Source Password
+
+    OAuth2 (Optional - instead of password):
+    SRC_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    SRC_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
+  Destination Account:
+    DEST_IMAP_HOST      : Destination IMAP Host
+    DEST_IMAP_USERNAME  : Destination Username/Email
+    DEST_IMAP_PASSWORD  : Destination Password
+
+    OAuth2 (Optional - instead of password):
+    DEST_OAUTH2_CLIENT_ID     : OAuth2 Client ID
+    DEST_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
+
+Also supports local folders as source and/or destination:
+    SRC_LOCAL_PATH      : Source local folder (backup root)
+    DEST_LOCAL_PATH     : Destination local folder (backup root)
+
+Usage:
+  python3 compare_imap_folders.py
+
+Examples:
+        # IMAP -> IMAP
+        python3 compare_imap_folders.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # Local -> IMAP
+        python3 compare_imap_folders.py \
+            --src-path "./my_backup" \
+            --dest-host "imap.dest.com" \
+            --dest-user "dest@example.com" \
+            --dest-pass "dest-app-password"
+
+        # IMAP -> Local
+        python3 compare_imap_folders.py \
+            --src-host "imap.source.com" \
+            --src-user "source@example.com" \
+            --src-pass "source-app-password" \
+            --dest-path "./my_backup"
 """
 
-import sys
+import argparse
+import os
+
+import imap_common
+import imap_oauth2
+
+
+def get_email_count(conn, folder_name):
+    """Return the IMAP message count for a folder, or None on error."""
+    try:
+        # Select folder in read-only mode
+        # Quote folder name handles spaces
+        typ, data = conn.select(f'"{folder_name}"', readonly=True)
+        if typ != "OK":
+            return None
+
+        # SELECT command returns the number of messages in data[0]
+        # data[0] is bytes, e.g. b'123'
+        if data and data[0]:
+            return int(data[0])
+        return 0
+
+    except Exception:
+        # print(f"Error checking {folder_name}: {e}")
+        return None
+
+
+def main():
+    default_src_path = os.getenv("SRC_LOCAL_PATH")
+    default_dest_path = os.getenv("DEST_LOCAL_PATH")
+
+    # Phase 1: determine whether each side is local
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--src-path", default=default_src_path)
+    pre_parser.add_argument("--dest-path", default=default_dest_path)
+    pre_args, _ = pre_parser.parse_known_args()
+
+    src_requires_imap = not bool(pre_args.src_path)
+    dest_requires_imap = not bool(pre_args.dest_path)
+
+    # Phase 2: full parser with conditional requirements
+    parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
+    parser.add_argument(
+        "--src-path",
+        default=default_src_path,
+        help="Source local folder (backup root). If set, IMAP source args are ignored.",
+    )
+    parser.add_argument(
+        "--dest-path",
+        default=default_dest_path,
+        help="Destination local folder (backup root). If set, IMAP destination args are ignored.",
+    )
+
+    # Source args
+    default_src_host = os.getenv("SRC_IMAP_HOST")
+    default_src_user = os.getenv("SRC_IMAP_USERNAME")
+    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
+    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--src-host",
+        default=default_src_host,
+        required=src_requires_imap and not bool(default_src_host),
+        help="Source IMAP Server (or SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--src-user",
+        default=default_src_user,
+        required=src_requires_imap and not bool(default_src_user),
+        help="Source Username (or SRC_IMAP_USERNAME)",
+    )
+    src_auth_required = src_requires_imap and not bool(default_src_pass or default_src_client_id)
+    src_auth = parser.add_mutually_exclusive_group(required=src_auth_required)
+    src_auth.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
+    src_auth.add_argument(
+        "--src-oauth2-client-id",
+        default=default_src_client_id,
+        dest="src_client_id",
+        help="Source OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
+    )
+    parser.add_argument(
+        "--src-oauth2-client-secret",
+        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="src_client_secret",
+        help="Source OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
+    )
+
+    # Dest args
+    default_dest_host = os.getenv("DEST_IMAP_HOST")
+    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
+    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
+    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--dest-host",
+        default=default_dest_host,
+        required=dest_requires_imap and not bool(default_dest_host),
+        help="Destination IMAP Server (or DEST_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--dest-user",
+        default=default_dest_user,
+        required=dest_requires_imap and not bool(default_dest_user),
+        help="Destination Username (or DEST_IMAP_USERNAME)",
+    )
+    dest_auth_required = dest_requires_imap and not bool(default_dest_pass or default_dest_client_id)
+    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
+    dest_auth.add_argument(
+        "--dest-pass",
+        default=default_dest_pass,
+        help="Destination Password (or DEST_IMAP_PASSWORD)",
+    )
+    dest_auth.add_argument(
+        "--dest-oauth2-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
+    )
+    dest_auth.add_argument(
+        "--dest-client-id",
+        default=default_dest_client_id,
+        dest="dest_client_id",
+        help=argparse.SUPPRESS,
+    )
+    parser.add_argument(
+        "--dest-oauth2-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--dest-client-secret",
+        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
+        dest="dest_client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    args = parser.parse_args()
+
+    src_is_local = bool(args.src_path)
+    dest_is_local = bool(args.dest_path)
+
+    SRC_HOST = args.src_host
+    SRC_USER = args.src_user
+    DEST_HOST = args.dest_host
+    DEST_USER = args.dest_user
+
+    # Acquire OAuth2 tokens if configured
+    src_oauth2_token = None
+    src_oauth2_provider = None
+    if not src_is_local and args.src_client_id:
+        src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
+            SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
+        )
+
+    dest_oauth2_token = None
+    dest_oauth2_provider = None
+    if not dest_is_local and args.dest_client_id:
+        dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
+            DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
+        )
+
+    print("\n--- Configuration Summary ---")
+    if src_is_local:
+        print(f"Source (Local)  : {args.src_path}")
+    else:
+        print(f"Source Host     : {args.src_host}")
+        print(f"Source User     : {args.src_user}")
+        print(f"Source Auth     : {imap_oauth2.auth_description(src_oauth2_provider)}")
+
+    if dest_is_local:
+        print(f"Destination (Local): {args.dest_path}")
+    else:
+        print(f"Destination Host: {args.dest_host}")
+        print(f"Destination User: {args.dest_user}")
+        print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
+    print("-----------------------------\n")
+
+    src = None
+    dest = None
+
+    try:
+        if not src_is_local:
+            # Connect to Source
+            print("Connecting to Source...")
+            src = imap_common.get_imap_connection(args.src_host, args.src_user, args.src_pass, src_oauth2_token)
+            if not src:
+                return
+
+        if not dest_is_local:
+            # Connect to Dest
+            print("Connecting to Destination...")
+            dest = imap_common.get_imap_connection(args.dest_host, args.dest_user, args.dest_pass, dest_oauth2_token)
+            if not dest:
+                return
+
+        # List Source Folders
+        print("Listing folders in Source...")
+        if src_is_local:
+            folders = imap_common.list_local_folders(args.src_path)
+        else:
+            folders = imap_common.list_selectable_folders(src)
+
+        if not folders:
+            print("Failed to list source folders.")
+            return
+
+        # Prepare Table Header
+        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
+        print("-" * len(header))
+        print(header)
+        print("-" * len(header))
+
+        total_src = 0
+        total_dest = 0
+
+        # Iterate through Source folders
+        for folder_name in folders:
+            # Get Counts
+            if src_is_local:
+                src_count = imap_common.get_local_email_count(args.src_path, folder_name)
+            else:
+                src_count = get_email_count(src, folder_name)
+
+            if dest_is_local:
+                dest_count = imap_common.get_local_email_count(args.dest_path, folder_name)
+            else:
+                dest_count = get_email_count(dest, folder_name)
+
+            # Format for display
+            src_str = str(src_count) if src_count is not None else "Err"
+            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
+
+            diff_str = ""
+            if src_count is not None and dest_count is not None:
+                diff = src_count - dest_count
+                diff_str = str(diff)
+                total_src += src_count
+                total_dest += dest_count
+            elif src_count is not None:
+                total_src += src_count
+
+            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
+
+        print("-" * len(header))
+        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
+
+    except Exception as e:
+        print(f"\nFatal Error: {e}")
+        if "Too many simultaneous connections" in str(e):
+            print("Tip: Wait a few minutes for previous connections to timeout or check other active scripts.")
+
+    finally:
+        # Check source connection state and logout if possible
+        if src:
+            try:
+                src.logout()
+            except Exception:
+                pass
+
+        # Check dest connection state and logout if possible
+        if dest:
+            try:
+                dest.logout()
+            except Exception:
+                pass
 
-from imap_compare import main
 
 if __name__ == "__main__":
-    print(
-        "WARNING: 'compare_imap_folders.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_compare.py'.",
-        file=sys.stderr,
-    )
-    print("Redirecting to 'imap_compare.py'...", file=sys.stderr)
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        sys.exit(1)
+    main()

src/count_imap_emails.py

@@ -1,24 +1,239 @@
-#!/usr/bin/env python3
-"""
-DEPRECATED: This script has been renamed.
-Please use imap_count.py instead.
+"""IMAP Email Counting Script.
+
+Counts emails per folder from either:
+- An IMAP account, or
+- A local backup folder created by ``backup_imap_emails.py`` (counts ``.eml`` files).
+
+Configuration (Environment Variables):
+    IMAP_HOST        : IMAP Host (e.g., imap.gmail.com)
+    IMAP_USERNAME    : Username/Email
+    IMAP_PASSWORD    : Password (or App Password)
+
+    OAuth2 (Optional - instead of password):
+    OAUTH2_CLIENT_ID        : OAuth2 Client ID
+    OAUTH2_CLIENT_SECRET    : OAuth2 Client Secret (required for Google)
+    SRC_OAUTH2_CLIENT_ID    : Alternate OAuth2 client ID env var
+    SRC_OAUTH2_CLIENT_SECRET: Alternate OAuth2 client secret env var
+
+Local backup counting:
+    BACKUP_LOCAL_PATH : Local backup root (preferred)
+    SRC_LOCAL_PATH    : Alternate local backup root
+
+Examples:
+    # Count an IMAP account
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export IMAP_PASSWORD="secretpassword"
+    python3 count_imap_emails.py
+
+    # Count an IMAP account using OAuth2
+    export IMAP_HOST="imap.gmail.com"
+    export IMAP_USERNAME="user@gmail.com"
+    export OAUTH2_CLIENT_ID="your-client-id"
+    export OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
+    python3 count_imap_emails.py
+
+    # Count a local backup
+    python3 count_imap_emails.py --path "./my_backup"
+
+    # Or set a default local backup path via env var
+    export BACKUP_LOCAL_PATH="./my_backup"
+    python3 count_imap_emails.py
 """
 
+import argparse
+import imaplib
+import os
 import sys
+from typing import Optional
+
+import imap_common
+import imap_oauth2
+
+
+def count_emails(imap_server, username, password=None, oauth2_token=None):
+    try:
+        # Connect to the IMAP server (using SSL)
+        print(f"Connecting to {imap_server}...")
+        mail = imap_common.get_imap_connection(imap_server, username, password, oauth2_token)
+        if not mail:
+            return
+
+        # List all mailboxes
+        print("Listing mailboxes...")
+        folders = imap_common.list_selectable_folders(mail)
+
+        if not folders:
+            print("Failed to list mailboxes.")
+            return
+
+        total_all_folders = 0
+        print(f"{'Folder Name':<40} {'Count':>10}")
+        print("-" * 52)
+
+        for folder_name in folders:
+            display_name = folder_name
+
+            try:
+                # Select the mailbox (read-only is sufficient for counting)
+                # folder_name extracted from list usually handles quotes correctly for select
+                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
+                if rv != "OK":
+                    print(f"{display_name:<40} {'Skipped':>10}")
+                    continue
+
+                # Search for all emails
+                status, data = mail.search(None, "ALL")
+
+                if status == "OK":
+                    # data[0] is space separated IDs
+                    email_ids = data[0].split()
+                    count = len(email_ids)
+                    print(f"{display_name:<40} {count:>10}")
+                    total_all_folders += count
+                else:
+                    print(f"{display_name:<40} {'Error':>10}")
+
+            except imaplib.IMAP4.error:
+                print(f"{display_name:<40} {'Error':>10}")
+
+        print("-" * 52)
+        print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+        # Logout
+        mail.logout()
+
+    except imaplib.IMAP4.error as e:
+        print(f"IMAP Error: {e}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def count_local_emails(local_path: str) -> None:
+    print(f"Scanning local backup: {local_path}")
+
+    folders = imap_common.list_local_folders(local_path)
+    if not folders:
+        print("No folders found.")
+        return
+
+    total_all_folders = 0
+    print(f"{'Folder Name':<40} {'Count':>10}")
+    print("-" * 52)
+
+    for folder_name in folders:
+        count = imap_common.get_local_email_count(local_path, folder_name)
+        if count is None:
+            print(f"{folder_name:<40} {'N/A':>10}")
+            continue
+
+        print(f"{folder_name:<40} {count:>10}")
+        total_all_folders += count
+
+    print("-" * 52)
+    print(f"{'TOTAL':<40} {total_all_folders:>10}")
+
+
+def main(argv: Optional[list[str]] = None) -> None:
+    # Phase 1: determine whether we're in local mode (--path)
+    default_path = os.getenv("BACKUP_LOCAL_PATH") or os.getenv("SRC_LOCAL_PATH")
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--path", default=default_path)
+    pre_args, _ = pre_parser.parse_known_args(argv)
+    require_imap = not bool(pre_args.path)
+
+    # Phase 2: full parser with conditional requirements
+    parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
+
+    parser.add_argument(
+        "--path",
+        default=default_path,
+        help="Local backup root to count (counts .eml files per folder). If set, IMAP args are ignored.",
+    )
+
+    # Try to unify var names for defaults. Priority: IMAP_* > SRC_IMAP_* > None
+    default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
+    default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
+    default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
+    default_client_id = os.getenv("OAUTH2_CLIENT_ID") or os.getenv("SRC_OAUTH2_CLIENT_ID")
+
+    parser.add_argument(
+        "--host",
+        default=default_host,
+        required=require_imap and not bool(default_host),
+        help="IMAP Server (or IMAP_HOST / SRC_IMAP_HOST)",
+    )
+    parser.add_argument(
+        "--user",
+        default=default_user,
+        required=require_imap and not bool(default_user),
+        help="Username (or IMAP_USERNAME / SRC_IMAP_USERNAME)",
+    )
+
+    auth_required = require_imap and not bool(default_pass or default_client_id)
+    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
+    auth_group.add_argument(
+        "--pass", dest="password", default=default_pass, help="Password (or IMAP_PASSWORD / SRC_IMAP_PASSWORD)"
+    )
+    auth_group.add_argument(
+        "--oauth2-client-id",
+        default=default_client_id,
+        dest="client_id",
+        help="OAuth2 Client ID (or OAUTH2_CLIENT_ID / SRC_OAUTH2_CLIENT_ID)",
+    )
+    auth_group.add_argument(
+        "--client-id",
+        default=default_client_id,
+        dest="client_id",
+        help=argparse.SUPPRESS,
+    )
+
+    parser.add_argument(
+        "--oauth2-client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help="OAuth2 Client Secret (if required) (or OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET)",
+    )
+    parser.add_argument(
+        "--client-secret",
+        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
+        dest="client_secret",
+        help=argparse.SUPPRESS,
+    )
+
+    args = parser.parse_args(argv)
+
+    if args.path:
+        if not os.path.isdir(args.path):
+            print(f"Error: Local path does not exist or is not a directory: {args.path}")
+            sys.exit(1)
+
+        print("\n--- Configuration Summary ---")
+        print(f"Local Path      : {args.path}")
+        print("-----------------------------\n")
+        count_local_emails(args.path)
+        raise SystemExit(0)
+
+    IMAP_SERVER = args.host
+    USERNAME = args.user
+    PASSWORD = args.password
+
+    # Acquire OAuth2 token if configured
+    oauth2_token = None
+    oauth2_provider = None
+    if args.client_id:
+        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
+            IMAP_SERVER, args.client_id, USERNAME, args.client_secret
+        )
+
+    print("\n--- Configuration Summary ---")
+    print(f"Host            : {IMAP_SERVER}")
+    print(f"User            : {USERNAME}")
+    print(f"Auth Method     : {imap_oauth2.auth_description(oauth2_provider)}")
+    print("-----------------------------\n")
+
+    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
 
-from imap_count import main
 
 if __name__ == "__main__":
-    print(
-        "WARNING: 'count_imap_emails.py' is deprecated and will be removed in a future version. Exact functionality is now available via 'imap_count.py'.",
-        file=sys.stderr,
-    )
-    print("Redirecting to 'imap_count.py'...", file=sys.stderr)
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        sys.exit(1)
+    main()

src/imap_backup.py

@@ -1,899 +0,0 @@
-"""
-IMAP Email Backup Script
-
-Backs up emails from an IMAP account to a local directory.
-Stores each email as a separate .eml file (RFC 5322 format) which is compatible with
-most email clients (Thunderbird, Apple Mail, Outlook, etc.).
-
-Features:
-- Incremental Backup: Skips messages that have already been downloaded (checks existing UIDs locally).
-- Filename Sanitization: Saves files as "{UID}_{Subject}.eml" with unsafe characters removed.
-- Folder Replication: Recreates the IMAP folder structure locally.
-- Parallel Processing: Uses multithreading for fast downloads.
-- Gmail Labels Preservation: Creates a manifest mapping Message-IDs to Gmail labels for restoration.
-
-Configuration (Environment Variables):
-    SRC_IMAP_HOST, SRC_IMAP_USERNAME: Source credentials.
-    SRC_IMAP_PASSWORD: Source password (or App Password).
-
-    OAuth2 (Optional - instead of password):
-    SRC_OAUTH2_CLIENT_ID: OAuth2 Client ID
-    SRC_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
-
-  BACKUP_LOCAL_PATH: Destination local directory.
-  MAX_WORKERS: Number of concurrent threads (default: 10).
-  BATCH_SIZE: Number of emails to process per batch (default: 10).
-  PRESERVE_LABELS: Set to "true" to create labels_manifest.json (Gmail). Default is "false".
-  PRESERVE_FLAGS: Set to "true" to preserve IMAP flags in manifest. Default is "false".
-  MANIFEST_ONLY: Set to "true" to only build manifest without downloading. Default is "false".
-  GMAIL_MODE: Set to "true" for Gmail backup mode. Default is "false".
-  DEST_DELETE: Set to "true" to delete local files not found on server (sync mode).
-              Default is "false".
-
-Usage:
-    python3 imap_backup.py \
-        --src-host "imap.example.com" \
-        --src-user "you@example.com" \
-        --src-pass "your-app-password" \
-        --dest-path "./my_backup"
-
-Gmail Labels:
-    python3 imap_backup.py \
-        --src-host "imap.gmail.com" \
-        --src-user "you@gmail.com" \
-        --src-pass "your-app-password" \
-        --dest-path "./my_backup" \
-        --preserve-labels \
-        "[Gmail]/All Mail"
-  This backs up all emails from [Gmail]/All Mail and creates a labels_manifest.json
-  file that maps each email's Message-ID to its Gmail labels for later restoration.
-"""
-
-import argparse
-import concurrent.futures
-import json
-import os
-import sys
-import threading
-
-from auth import imap_oauth2
-from core import imap_session
-from providers import provider_exchange, provider_gmail
-from utils import imap_common
-
-# Defaults
-MAX_WORKERS = 10
-BATCH_SIZE = 10
-MANIFEST_FILENAME = "labels_manifest.json"
-
-# Thread-local storage
-thread_local = threading.local()
-safe_print = imap_common.safe_print
-
-
-def process_single_uid(src, uid, folder_name, local_folder_path):
-    """
-    Fetch and save a single email by UID.
-
-    Args:
-        src: IMAP connection
-        uid: Message UID to fetch
-        folder_name: Current folder name (for logging)
-        local_folder_path: Local directory to save email
-
-    Returns:
-        Tuple of (success: bool, connection):
-        - (True, src): UID processed successfully (or skipped)
-        - (False, src): Auth error occurred, caller should retry after reconnect
-    """
-    uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-
-    try:
-        resp, data = src.uid("fetch", uid, "(RFC822)")
-        if resp != "OK" or not data or data[0] is None:
-            safe_print(f"[{folder_name}] ERROR Fetch Body | UID {uid_str}")
-            return (True, src)  # Don't retry, move on
-
-        raw_email = None
-        for item in data:
-            if isinstance(item, tuple):
-                raw_email = item[1]
-                break
-
-        # Derive Subject for filename from the already-fetched message bytes.
-        _, subject = imap_common.parse_message_id_and_subject_from_bytes(raw_email)
-        if not subject:
-            clean_subject = "No Subject"
-        else:
-            clean_subject = imap_common.sanitize_filename(subject)
-            clean_subject = clean_subject[:100]
-
-        filename = f"{uid_str}_{clean_subject}.eml"
-        full_path = os.path.join(local_folder_path, filename)
-
-        if os.path.exists(full_path):
-            return (True, src)  # Already exists
-
-        if raw_email:
-            try:
-                with open(full_path, "wb") as f:
-                    f.write(raw_email)
-                safe_print(f"[{folder_name}] SAVED  | {filename[:60]}...")
-            except OSError as e:
-                safe_print(f"[{folder_name}] ERROR Write | {uid_str}: {e}")
-        else:
-            safe_print(f"[{folder_name}] EMPTY Content | UID {uid_str}")
-
-        return (True, src)
-
-    except Exception as e:
-        if imap_oauth2.is_auth_error(e):
-            safe_print(f"[{folder_name}] Auth error for UID {uid_str}, will retry...")
-            return (False, src)  # Signal retry needed
-        else:
-            safe_print(f"[{folder_name}] ERROR Processing UID {uid}: {e}")
-            return (True, src)  # Don't retry other errors
-
-
-def process_batch(uids, folder_name, src_conf, local_folder_path):
-    src = imap_session.get_thread_connection(thread_local, "src", src_conf)
-    if not src:
-        safe_print("Error: Could not establish connection for batch.")
-        return
-
-    try:
-        src.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Error selecting folder {folder_name} in worker: {e}")
-        return
-
-    for uid in uids:
-        uid_str = uid.decode("utf-8") if isinstance(uid, bytes) else str(uid)
-        max_retries = 2
-
-        for attempt in range(max_retries):
-            src, ok = imap_session.ensure_folder_session(src, src_conf, folder_name, readonly=True)
-            thread_local.src = src
-            if not ok:
-                safe_print(f"[{folder_name}] ERROR: Connection/folder lost for UID {uid_str}")
-                return
-
-            success, src = process_single_uid(src, uid, folder_name, local_folder_path)
-            thread_local.src = src
-
-            if success:
-                break
-            if attempt < max_retries - 1:
-                src = None
-                thread_local.src = None
-
-
-def get_existing_uids(local_path):
-    """
-    Scans the local directory for files matching pattern matches {UID}_*.eml
-    Returns a set of UIDs (as strings).
-    """
-    existing = set()
-    if not os.path.exists(local_path):
-        return existing
-
-    try:
-        for filename in os.listdir(local_path):
-            if filename.endswith(".eml") and "_" in filename:
-                # Expecting UID_Subject.eml
-                parts = filename.split("_", 1)
-                if parts[0].isdigit():
-                    existing.add(parts[0])
-    except Exception:
-        pass
-    return existing
-
-
-# Standard IMAP flags that can be preserved during migration
-# \Recent is session-specific and cannot be set by clients
-# \Deleted should not be preserved as it marks messages for removal
-PRESERVABLE_FLAGS = imap_common.PRESERVABLE_FLAGS
-
-
-def get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
-    """
-    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder,
-    with OAuth2 session management.
-
-    Args:
-        imap_conn: IMAP connection
-        folder_name: Folder to scan
-        src_conf: Connection config dict for OAuth2 refresh
-        progress_callback: Optional callback(current, total) for progress reporting
-
-    Returns:
-        Tuple of (message_info, imap_conn) where message_info is:
-        { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
-        The returned imap_conn may be different if reconnection occurred.
-    """
-    message_info = {}
-
-    # Ensure connection is healthy (refresh OAuth2 token if needed) before initial select
-    if src_conf:
-        imap_conn = imap_session.ensure_connection(imap_conn, src_conf)
-        if not imap_conn:
-            safe_print(f"Could not establish connection for folder {folder_name}")
-            return (message_info, None)
-
-    try:
-        imap_conn.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Could not select folder {folder_name}: {e}")
-        return (message_info, imap_conn)
-
-    try:
-        resp, data = imap_conn.uid("search", None, "ALL")
-        if resp != "OK" or not data or not data[0]:
-            return (message_info, imap_conn)
-
-        uids = data[0].split()
-        if not uids:
-            return (message_info, imap_conn)
-
-        total_uids = len(uids)
-
-        # Fetch Message-IDs and FLAGS in batches - use larger batch for header-only fetches
-        batch_size = 200
-        for i in range(0, len(uids), batch_size):
-            batch = uids[i : i + batch_size]
-            uid_range = b",".join(batch)
-
-            # Report progress
-            if progress_callback:
-                progress_callback(min(i + batch_size, total_uids), total_uids)
-
-            # Proactively refresh token and ensure folder is selected
-            if src_conf:
-                imap_conn, ok = imap_session.ensure_folder_session(imap_conn, src_conf, folder_name, readonly=True)
-                if not ok:
-                    safe_print(f"ERROR: Connection/folder lost in {folder_name}")
-                    break
-
-            try:
-                # Fetch both Message-ID header and FLAGS
-                resp, items = imap_conn.uid("fetch", uid_range, "(FLAGS BODY.PEEK[HEADER.FIELDS (MESSAGE-ID)])")
-                if resp != "OK":
-                    continue
-
-                # Parse response - items come in pairs for each message
-                for item in items:
-                    if isinstance(item, tuple) and len(item) >= 2:
-                        # First element contains UID and FLAGS info
-                        meta_str = (
-                            item[0].decode("utf-8", errors="ignore") if isinstance(item[0], bytes) else str(item[0])
-                        )
-
-                        # Extract all preservable flags from the metadata
-                        flags = []
-                        for flag in imap_common.PRESERVABLE_FLAGS:
-                            if flag in meta_str:
-                                flags.append(flag)
-
-                        # Second element contains the header
-                        msg_id = imap_common.extract_message_id(item[1])
-                        if msg_id:
-                            message_info[msg_id] = {"flags": flags}
-            except Exception as e:
-                safe_print(f"Error fetching batch in {folder_name}: {e}")
-                continue
-
-    except Exception as e:
-        safe_print(f"Error searching folder {folder_name}: {e}")
-
-    return (message_info, imap_conn)
-
-
-def get_message_info_in_folder(imap_conn, folder_name, progress_callback=None):
-    """
-    Returns a dict of Message-IDs and their IMAP flags for all emails in a folder.
-
-    Returns: { "message-id": {"flags": ["\\Seen", "\\Flagged", ...]}, ... }
-    Optional progress_callback(current, total) for progress reporting.
-    """
-    message_info, _ = get_message_info_in_folder_with_conf(imap_conn, folder_name, None, progress_callback)
-    return message_info
-
-
-def get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback=None):
-    """
-    Returns a set of Message-IDs for all emails in a given folder,
-    with OAuth2 session management.
-
-    Args:
-        imap_conn: IMAP connection
-        folder_name: Folder to scan
-        src_conf: Connection config dict for OAuth2 refresh
-        progress_callback: Optional callback(current, total) for progress reporting
-
-    Returns:
-        Tuple of (message_ids, imap_conn) where message_ids is a set of strings.
-        The returned imap_conn may be different if reconnection occurred.
-    """
-    info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_callback)
-    return (set(info.keys()), imap_conn)
-
-
-def get_message_ids_in_folder(imap_conn, folder_name, progress_callback=None):
-    """
-    Returns a set of Message-IDs for all emails in a given folder.
-    This is a convenience wrapper around get_message_info_in_folder.
-    Optional progress_callback(current, total) for progress reporting.
-    """
-    info = get_message_info_in_folder(imap_conn, folder_name, progress_callback)
-    return set(info.keys())
-
-
-def build_labels_manifest(imap_conn, local_path, src_conf=None):
-    """
-    Builds a manifest mapping Message-IDs to their Gmail labels and IMAP flags.
-    Scans all folders (labels) in the account and records which Message-IDs
-    appear in each label, plus their flags from [Gmail]/All Mail.
-
-    Returns a dict: { "message-id": {"labels": ["Label1", ...], "flags": ["\\Seen", ...]}, ... }
-    Saves the manifest to labels_manifest.json in the backup directory.
-    Optional src_conf for automatic token refresh on expiration.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-    total_emails_scanned = 0
-
-    safe_print("--- Building Gmail Labels Manifest ---")
-
-    # Get all selectable folders and filter to label folders
-    all_folders = imap_common.list_selectable_folders(imap_conn)
-    if not all_folders:
-        safe_print("Error: Could not list folders for label mapping.")
-        return manifest
-
-    # First, scan [Gmail]/All Mail to get the authoritative flags for all emails
-    safe_print("[0/N] Scanning [Gmail]/All Mail for flags (read/starred/etc)...")
-    all_mail_start = time.time()
-
-    def all_mail_progress_cb(current, total):
-        elapsed = time.time() - start_time
-        print(
-            f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-            end="",
-            flush=True,
-        )
-
-    all_mail_info, imap_conn = get_message_info_in_folder_with_conf(
-        imap_conn, provider_gmail.GMAIL_ALL_MAIL, src_conf, all_mail_progress_cb
-    )
-    print()  # New line after progress
-
-    # Initialize manifest with flags from All Mail
-    for msg_id, info in all_mail_info.items():
-        manifest[msg_id] = {"labels": [], "flags": info.get("flags", [])}
-
-    all_mail_elapsed = time.time() - all_mail_start
-    # Count flag statistics
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print(f"  -> {len(all_mail_info)} emails scanned ({all_mail_elapsed:.1f}s)")
-    safe_print(f"  -> Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}\n")
-
-    # Keep connection alive
-    try:
-        imap_conn.noop()
-    except Exception:
-        pass
-
-    # Parse folder names and filter to label folders
-    label_folders = [f for f in all_folders if provider_gmail.is_label_folder(f)]
-
-    total_folders = len(label_folders)
-    safe_print(f"Found {total_folders} label folders to scan.\n")
-
-    # Scan each label folder
-    for folder_idx, folder_name in enumerate(label_folders, 1):
-        folder_start = time.time()
-
-        # Progress callback for this folder
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        message_ids, imap_conn = get_message_ids_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
-        print()  # New line after progress
-
-        # Keep connection alive between folders
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-        # Determine the label name to store
-        # For [Gmail]/Sent Mail -> "Sent Mail"
-        # For [Gmail]/Starred -> "Starred"
-        # For INBOX -> "INBOX"
-        # For user folders -> folder name as-is
-        if folder_name.startswith("[Gmail]/"):
-            label_name = folder_name[8:]  # Remove "[Gmail]/" prefix
-        else:
-            label_name = folder_name
-
-        for msg_id in message_ids:
-            if msg_id not in manifest:
-                # Email not in All Mail (rare, but handle it)
-                manifest[msg_id] = {"labels": [], "flags": []}
-            if label_name not in manifest[msg_id]["labels"]:
-                manifest[msg_id]["labels"].append(label_name)
-
-        folder_elapsed = time.time() - folder_start
-        total_emails_scanned += len(message_ids)
-        safe_print(f"  -> {len(message_ids)} emails with label '{label_name}' ({folder_elapsed:.1f}s)")
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print("\nManifest building complete:")
-    safe_print(f"  - Folders scanned: {total_folders + 1}")  # +1 for All Mail
-    safe_print(f"  - Total email-label mappings: {total_emails_scanned}")
-    safe_print(f"  - Unique emails: {len(manifest)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Starred: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
-    try:
-        with open(manifest_path, "w", encoding="utf-8") as f:
-            json.dump(manifest, f, indent=2, ensure_ascii=False)
-        safe_print(f"\nLabels manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def build_flags_manifest(imap_conn, local_path, folders_to_scan=None, src_conf=None):
-    """
-    Builds a manifest mapping Message-IDs to their IMAP flags.
-    For non-Gmail servers, scans specified folders (or all folders if not specified).
-
-    Returns a dict: { "message-id": {"flags": ["\\Seen", ...]}, ... }
-    Saves the manifest to flags_manifest.json in the backup directory.
-    Optional src_conf for automatic token refresh on expiration.
-    """
-    import time
-
-    manifest = {}
-    start_time = time.time()
-
-    safe_print("--- Building Flags Manifest ---")
-
-    # Get folders to scan
-    if folders_to_scan:
-        all_folders = folders_to_scan
-    else:
-        try:
-            typ, folders = imap_conn.list()
-            if typ != "OK":
-                safe_print("Error: Could not list folders.")
-                return manifest
-            all_folders = [imap_common.normalize_folder_name(f) for f in folders]
-        except Exception as e:
-            safe_print(f"Error listing folders: {e}")
-            return manifest
-
-    total_folders = len(all_folders)
-    safe_print(f"Found {total_folders} folders to scan.\n")
-
-    # Scan each folder
-    for folder_idx, folder_name in enumerate(all_folders, 1):
-        folder_start = time.time()
-
-        def progress_cb(current, total):
-            elapsed = time.time() - start_time
-            print(
-                f"\r  Progress: {current}/{total} emails scanned (elapsed: {elapsed:.0f}s)   ",
-                end="",
-                flush=True,
-            )
-
-        safe_print(f"[{folder_idx}/{total_folders}] Scanning: {folder_name}")
-        folder_info, imap_conn = get_message_info_in_folder_with_conf(imap_conn, folder_name, src_conf, progress_cb)
-        print()  # New line after progress
-
-        # Merge into manifest (keep first occurrence of flags)
-        for msg_id, info in folder_info.items():
-            if msg_id not in manifest:
-                manifest[msg_id] = {"flags": info.get("flags", [])}
-
-        folder_elapsed = time.time() - folder_start
-        safe_print(f"  -> {len(folder_info)} emails ({folder_elapsed:.1f}s)")
-
-        # Keep connection alive
-        try:
-            imap_conn.noop()
-        except Exception:
-            pass
-
-    # Summary
-    total_elapsed = time.time() - start_time
-    read_count = sum(1 for m in manifest.values() if imap_common.FLAG_SEEN in m.get("flags", []))
-    flagged_count = sum(1 for m in manifest.values() if imap_common.FLAG_FLAGGED in m.get("flags", []))
-    safe_print("\nFlags manifest building complete:")
-    safe_print(f"  - Folders scanned: {total_folders}")
-    safe_print(f"  - Unique emails: {len(manifest)}")
-    safe_print(f"  - Read: {read_count}, Unread: {len(manifest) - read_count}, Flagged: {flagged_count}")
-    safe_print(f"  - Time elapsed: {total_elapsed:.1f}s")
-
-    # Save manifest
-    manifest_path = os.path.join(local_path, "flags_manifest.json")
-    try:
-        with open(manifest_path, "w", encoding="utf-8") as f:
-            json.dump(manifest, f, indent=2, ensure_ascii=False)
-        safe_print(f"\nFlags manifest saved to: {manifest_path}")
-    except Exception as e:
-        safe_print(f"Error saving manifest: {e}")
-
-    return manifest
-
-
-def delete_orphan_local_files(local_folder_path, server_uids):
-    """
-    Delete local .eml files that no longer exist on the server.
-    Args:
-        local_folder_path: Path to local folder containing .eml files
-        server_uids: Set of UID strings currently on server
-    Returns:
-        Count of deleted files
-    """
-    deleted_count = 0
-    if not os.path.exists(local_folder_path):
-        return deleted_count
-
-    try:
-        for filename in os.listdir(local_folder_path):
-            if not filename.endswith(".eml") or "_" not in filename:
-                continue
-
-            # Extract UID from filename (format: {UID}_{Subject}.eml)
-            parts = filename.split("_", 1)
-            if not parts[0].isdigit():
-                continue
-
-            local_uid = parts[0]
-            if local_uid not in server_uids:
-                file_path = os.path.join(local_folder_path, filename)
-                try:
-                    os.remove(file_path)
-                    safe_print(f"  -> Deleted orphan: {filename}")
-                    deleted_count += 1
-                except Exception as e:
-                    safe_print(f"  -> Error deleting {filename}: {e}")
-    except Exception as e:
-        safe_print(f"Error scanning for orphan files: {e}")
-
-    return deleted_count
-
-
-def backup_folder(src_main, folder_name, local_base_path, src_conf, dest_delete=False):
-    safe_print(f"--- Processing Folder: {folder_name} ---")
-
-    # create local path
-    # Handle folder separators. IMAP output might be "Parent/Child"
-    # We rely on OS to handle "Parent/Child" as subdirectories using join
-    # But clean the segments
-    cleaned_name = folder_name.replace("/", os.sep)
-    local_folder_path = os.path.join(local_base_path, cleaned_name)
-
-    try:
-        os.makedirs(local_folder_path, exist_ok=True)
-    except Exception as e:
-        safe_print(f"Error creating directory {local_folder_path}: {e}")
-        return
-
-    # Select IMAP folder
-    try:
-        src_main.select(f'"{folder_name}"', readonly=True)
-    except Exception as e:
-        safe_print(f"Skipping {folder_name}: {e}")
-        return
-
-    # Search all
-    resp, data = src_main.uid("search", None, "ALL")
-    if resp != "OK":
-        return
-
-    uids = data[0].split()
-    total_on_server = len(uids)
-
-    # Build set of server UIDs for comparison
-    server_uid_set = set()
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        server_uid_set.add(u_str)
-
-    if total_on_server == 0:
-        safe_print(f"Folder {folder_name} is empty.")
-        # If dest_delete enabled, delete all local files
-        if dest_delete:
-            deleted = delete_orphan_local_files(local_folder_path, set())
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-        return
-
-    # Incremental Optimization
-    # Read local directory to find UIDs we already have
-    existing_uids = get_existing_uids(local_folder_path)
-
-    # Delete orphan local files if dest_delete is enabled
-    if dest_delete:
-        orphan_uids = existing_uids - server_uid_set
-        if orphan_uids:
-            safe_print(f"Found {len(orphan_uids)} local files not on server, deleting...")
-            deleted = delete_orphan_local_files(local_folder_path, server_uid_set)
-            if deleted > 0:
-                safe_print(f"Deleted {deleted} orphan files from local backup.")
-
-    # Filter UIDs
-    # decode uid first if bytes
-    uids_to_download = []
-    for u in uids:
-        u_str = u.decode("utf-8") if isinstance(u, bytes) else str(u)
-        if u_str not in existing_uids:
-            uids_to_download.append(u)
-
-    skipped = total_on_server - len(uids_to_download)
-    if skipped > 0:
-        safe_print(f"Skipping {skipped} emails (already exist locally).")
-
-    if not uids_to_download:
-        safe_print(f"Folder {folder_name} is up to date.")
-        return
-
-    safe_print(f"Downloading {len(uids_to_download)} new emails...")
-
-    # Create batches
-    uid_batches = [uids_to_download[i : i + BATCH_SIZE] for i in range(0, len(uids_to_download), BATCH_SIZE)]
-
-    executor = concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)
-    try:
-        futures = []
-        for batch in uid_batches:
-            futures.append(executor.submit(process_batch, batch, folder_name, src_conf, local_folder_path))
-
-        for future in concurrent.futures.as_completed(futures):
-            future.result()
-
-    except KeyboardInterrupt:
-        executor.shutdown(wait=False, cancel_futures=True)
-        raise
-    finally:
-        executor.shutdown(wait=True)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Backup IMAP emails to local .eml files.")
-    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
-
-    # Source
-    default_src_host = os.getenv("SRC_IMAP_HOST")
-    default_src_user = os.getenv("SRC_IMAP_USERNAME")
-    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
-    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--src-host",
-        default=default_src_host,
-        required=not bool(default_src_host),
-        help="Source IMAP Server (or SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--src-user",
-        default=default_src_user,
-        required=not bool(default_src_user),
-        help="Source Username (or SRC_IMAP_USERNAME)",
-    )
-
-    # Authentication: require either password OR OAuth2 client-id (unless provided via env vars)
-    auth_required = not bool(default_src_pass or default_src_client_id)
-    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
-    auth_group.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
-    # OAuth2
-    auth_group.add_argument(
-        "--src-oauth2-client-id",
-        default=default_src_client_id,
-        dest="src_client_id",
-        help="OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
-    )
-    parser.add_argument(
-        "--src-oauth2-client-secret",
-        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="src_client_secret",
-        help="OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
-    )
-
-    # Destination (Local Path)
-    env_path = os.getenv("BACKUP_LOCAL_PATH")
-    parser.add_argument(
-        "--dest-path",
-        default=env_path,
-        required=not bool(env_path),
-        help="Local destination path (or BACKUP_LOCAL_PATH)",
-    )
-
-    # Config
-    parser.add_argument("--workers", type=int, default=int(os.getenv("MAX_WORKERS", 10)), help="Thread count")
-    parser.add_argument("--batch", type=int, default=int(os.getenv("BATCH_SIZE", 10)), help="Emails per batch")
-
-    # Gmail Labels
-    env_preserve_labels = os.getenv("PRESERVE_LABELS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-labels",
-        action="store_true",
-        default=env_preserve_labels,
-        help="Gmail only: Create a labels_manifest.json mapping Message-IDs to labels for restoration",
-    )
-    env_preserve_flags = os.getenv("PRESERVE_FLAGS", "false").lower() == "true"
-    parser.add_argument(
-        "--preserve-flags",
-        action="store_true",
-        default=env_preserve_flags,
-        help="Preserve IMAP flags (read/unread, starred, answered, draft) in manifest for restoration",
-    )
-    env_manifest_only = os.getenv("MANIFEST_ONLY", "false").lower() == "true"
-    parser.add_argument(
-        "--manifest-only",
-        action="store_true",
-        default=env_manifest_only,
-        help="Gmail only: Build the labels manifest and exit without downloading emails",
-    )
-    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
-    parser.add_argument(
-        "--gmail-mode",
-        action="store_true",
-        default=env_gmail_mode,
-        help="Gmail backup mode: Build labels manifest and backup [Gmail]/All Mail only (recommended)",
-    )
-
-    # Sync mode: delete local files not on server
-    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
-    parser.add_argument(
-        "--dest-delete",
-        action="store_true",
-        default=env_dest_delete,
-        help="Delete local .eml files that no longer exist on the IMAP server (sync mode)",
-    )
-
-    parser.add_argument("folder", nargs="?", help="Specific folder to backup")
-
-    args = parser.parse_args()
-
-    global MAX_WORKERS, BATCH_SIZE
-    MAX_WORKERS = args.workers
-    BATCH_SIZE = args.batch
-
-    # Build connection config (acquires OAuth2 token if configured)
-    src_conf = imap_session.build_imap_conf(
-        args.src_host, args.src_user, args.src_pass, args.src_client_id, args.src_client_secret
-    )
-
-    # Expand path (~/...)
-    local_path = os.path.expanduser(args.dest_path)
-    if not os.path.exists(local_path):
-        try:
-            os.makedirs(local_path)
-            print(f"Created backup directory: {local_path}")
-        except Exception as e:
-            print(f"Error creating backup directory: {e}")
-            sys.exit(1)
-
-    print("\n--- Configuration Summary ---")
-    print(f"Source Host     : {args.src_host}")
-    print(f"Source User     : {args.src_user}")
-    print(f"Auth Method     : {imap_oauth2.auth_description(src_conf['oauth2'] and src_conf['oauth2']['provider'])}")
-    print(f"Destination Path: {local_path}")
-    if args.gmail_mode:
-        print("Mode            : Gmail Backup (All Mail + Labels + Flags)")
-    elif args.manifest_only:
-        print("Mode            : Manifest Only (no email download)")
-    elif args.folder:
-        print(f"Target Folder   : {args.folder}")
-    if args.preserve_labels or args.manifest_only or args.gmail_mode:
-        print("Preserve Labels : Yes (Gmail)")
-    if args.preserve_flags or args.gmail_mode:
-        print("Preserve Flags  : Yes (read/starred/answered/draft)")
-    if args.dest_delete:
-        print("Dest Delete     : Yes (remove local orphans)")
-    print("-----------------------------\n")
-
-    try:
-        src = imap_common.get_imap_connection_from_conf(src_conf)
-        if not src:
-            sys.exit(1)
-
-        # Build labels manifest BEFORE backing up emails (Gmail mode)
-        # This way we capture the label state at backup time
-        if args.preserve_labels or args.manifest_only or args.gmail_mode:
-            print("Building Gmail labels manifest...")
-            print("This scans all folders to map Message-IDs to labels and flags.\n")
-            build_labels_manifest(src, local_path, src_conf)
-            print("")  # Blank line after manifest building
-        # Build flags-only manifest for non-Gmail servers
-        elif args.preserve_flags:
-            print("Building flags manifest...")
-            print("This scans folders to capture read/starred/etc status.\n")
-            # Get folders to scan
-            folders_to_scan = [args.folder] if args.folder else None
-            build_flags_manifest(src, local_path, folders_to_scan, src_conf)
-            print("")  # Blank line after manifest building
-
-        # If manifest-only mode, we're done
-        if args.manifest_only:
-            try:
-                src.logout()
-            except Exception:
-                pass  # Connection may already be closed
-            manifest_path = os.path.join(local_path, MANIFEST_FILENAME)
-            print("\nManifest-only mode complete.")
-            print(f"Labels manifest saved to: {manifest_path}")
-            print("\nTo download emails, run again without --manifest-only:")
-            print(f'  python3 imap_backup.py --dest-path "{local_path}" "[Gmail]/All Mail"')
-            sys.exit(0)
-
-        # Gmail mode: backup only [Gmail]/All Mail
-        if args.gmail_mode:
-            backup_folder(src, provider_gmail.GMAIL_ALL_MAIL, local_path, src_conf, args.dest_delete)
-        elif args.folder:
-            backup_folder(src, args.folder, local_path, src_conf, args.dest_delete)
-        else:
-            # Reconnect after potentially long manifest building
-            src = imap_session.ensure_connection(src, src_conf)
-            if not src:
-                print("Warning: Could not reconnect to IMAP server for backup. Manifest was saved successfully.")
-                sys.exit(0)
-            folders = imap_common.list_selectable_folders(src)
-            for name in folders:
-                if provider_exchange.is_special_folder(name):
-                    print(f"Skipping Exchange system folder: {name}")
-                    continue
-                src = imap_session.ensure_connection(src, src_conf)
-                if not src:
-                    print("Fatal: Could not reconnect to IMAP server. Aborting.")
-                    sys.exit(1)
-                backup_folder(src, name, local_path, src_conf, args.dest_delete)
-
-        try:
-            src.logout()
-        except Exception:
-            pass  # Connection may already be closed
-        print("\nBackup completed successfully.")
-
-        if args.preserve_labels or args.gmail_mode:
-            manifest_path = os.path.join(local_path, "labels_manifest.json")
-            print(f"\nGmail labels manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply labels and flags to emails.")
-        elif args.preserve_flags:
-            manifest_path = os.path.join(local_path, "flags_manifest.json")
-            print(f"\nFlags manifest saved to: {manifest_path}")
-            print("Use this file when restoring to reapply read/starred status to emails.")
-
-    except KeyboardInterrupt:
-        raise
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        sys.exit(1)

src/imap_common.py

@@ -16,39 +16,10 @@ import urllib.parse
 from email import policy
 from email.header import decode_header
 from email.parser import BytesParser
-from importlib.metadata import PackageNotFoundError, version
-
-from auth import imap_oauth2
-from utils import imap_compress, imap_retry, restore_cache
-
-
-def get_version() -> str:
-    """
-    Return the package version string.
-
-    Tries to get the version from the installed package metadata.
-    If not installed, falls back to reading pyproject.toml from the project root.
-    """
-    try:
-        return version("imap-migration-tools")
-    except PackageNotFoundError:
-        # Fallback: Try to read from pyproject.toml for local development
-        try:
-            # Assuming src/imap_common.py structure
-            root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-            pyproject_path = os.path.join(root_dir, "pyproject.toml")
-
-            if os.path.exists(pyproject_path):
-                with open(pyproject_path, encoding="utf-8") as f:
-                    for line in f:
-                        # Looking for: version = "1.2.3"
-                        match = re.search(r'^version\s*=\s*"([^"]+)"', line.strip())
-                        if match:
-                            return match.group(1)
-        except Exception:
-            pass
-        return "0.0.0-dev"
 
+import imap_oauth2
+import imap_retry
+import restore_cache
 
 # Standard IMAP flags
 FLAG_SEEN = "\\Seen"
@@ -319,7 +290,6 @@ def get_imap_connection(host, user, password=None, oauth2_token=None):
             conn.authenticate("XOAUTH2", lambda _: auth_string.encode())
         else:
             conn.login(user, password)
-        imap_compress.enable_compression(conn, log_fn=safe_print)
         return imap_retry.ConnectionProxy(conn, log_fn=safe_print)
     except Exception as e:
         print(f"Connection error to {host}: {e}")
@@ -493,6 +463,28 @@ def parse_message_id_and_subject_from_bytes(raw_message):
         return None, "(No Subject)"
 
 
+def message_exists_in_folder(dest_conn, msg_id):
+    """
+    Checks if a message with the given Message-ID exists in the CURRENTLY SELECTED folder of dest_conn.
+    Only considers non-deleted messages (UNDELETED) so that messages pending expunge
+    don't block uploads of fresh copies.
+    Returns True if found, False otherwise.
+    """
+    if not msg_id:
+        return False
+
+    clean_id = msg_id.replace('"', '\\"')
+    try:
+        typ, data = dest_conn.search(None, f'UNDELETED (HEADER Message-ID "{clean_id}")')
+        if typ != "OK":
+            return False
+
+        dest_ids = data[0].split()
+        return len(dest_ids) > 0
+    except Exception:
+        return False
+
+
 def get_message_ids_in_folder(imap_conn):
     """
     Fetches all Message-IDs from the currently selected folder.
@@ -714,55 +706,6 @@ def delete_orphan_emails(imap_conn, folder_name, source_msg_ids, dest_uid_to_msg
     return deleted_count
 
 
-def load_folder_msg_ids(
-    imap_conn,
-    folder_name,
-    msg_ids_by_folder,
-    msg_ids_lock,
-    progress_cache_data=None,
-    progress_cache_lock=None,
-    dest_host=None,
-    dest_user=None,
-):
-    """Load Message-IDs for a folder, fetching from server if not yet cached.
-
-    On first call for a folder, SELECTs the folder and fetches all Message-IDs
-    from the server (merged with progress cache). Subsequent calls return the
-    cached set without any IMAP operations.
-
-    Returns the set of Message-IDs, or None if tracking is disabled.
-    """
-    if msg_ids_by_folder is None or msg_ids_lock is None:
-        return None
-
-    with msg_ids_lock:
-        existing = msg_ids_by_folder.get(folder_name)
-
-    if existing is not None:
-        return existing
-
-    # Build from progress cache
-    built: set[str] = set()
-    if is_progress_cache_ready(progress_cache_data, progress_cache_lock) and dest_host and dest_user:
-        built = restore_cache.get_cached_message_ids(
-            progress_cache_data,
-            progress_cache_lock,
-            dest_host,
-            dest_user,
-            folder_name,
-        )
-
-    # Fetch from server for a comprehensive set (one-time cost per folder)
-    ensure_folder_exists(imap_conn, folder_name)
-    imap_conn.select(f'"{folder_name}"')
-    server_ids = set(get_message_ids_in_folder(imap_conn).values())
-    built.update(server_ids)
-
-    with msg_ids_lock:
-        msg_ids_by_folder.setdefault(folder_name, built)
-        return msg_ids_by_folder[folder_name]
-
-
 def load_manifest(local_path, filename):
     """Load a manifest JSON file from a backup directory.
 

src/imap_compare.py

@@ -1,337 +0,0 @@
-"""
-IMAP Folder Comparison Script
-
-This script compares email counts between a source and a destination.
-Each side can be either an IMAP account or a local backup folder.
-It iterates through all folders found in the source account and checks the corresponding
-folder in the destination account.
-
-Configuration (Environment Variables):
-  Source Account:
-    SRC_IMAP_HOST       : Source IMAP Host
-    SRC_IMAP_USERNAME   : Source Username/Email
-    SRC_IMAP_PASSWORD   : Source Password
-
-    OAuth2 (Optional - instead of password):
-    SRC_OAUTH2_CLIENT_ID     : OAuth2 Client ID
-    SRC_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
-
-  Destination Account:
-    DEST_IMAP_HOST      : Destination IMAP Host
-    DEST_IMAP_USERNAME  : Destination Username/Email
-    DEST_IMAP_PASSWORD  : Destination Password
-
-    OAuth2 (Optional - instead of password):
-    DEST_OAUTH2_CLIENT_ID     : OAuth2 Client ID
-    DEST_OAUTH2_CLIENT_SECRET : OAuth2 Client Secret (required for Google)
-
-Also supports local folders as source and/or destination:
-    SRC_LOCAL_PATH      : Source local folder (backup root)
-    DEST_LOCAL_PATH     : Destination local folder (backup root)
-
-Usage:
-  python3 imap_compare.py
-
-Examples:
-        # IMAP -> IMAP
-        python3 imap_compare.py \
-            --src-host "imap.source.com" \
-            --src-user "source@example.com" \
-            --src-pass "source-app-password" \
-            --dest-host "imap.dest.com" \
-            --dest-user "dest@example.com" \
-            --dest-pass "dest-app-password"
-
-        # Local -> IMAP
-        python3 imap_compare.py \
-            --src-path "./my_backup" \
-            --dest-host "imap.dest.com" \
-            --dest-user "dest@example.com" \
-            --dest-pass "dest-app-password"
-
-        # IMAP -> Local
-        python3 imap_compare.py \
-            --src-host "imap.source.com" \
-            --src-user "source@example.com" \
-            --src-pass "source-app-password" \
-            --dest-path "./my_backup"
-"""
-
-import argparse
-import os
-import sys
-
-from auth import imap_oauth2
-from utils import imap_common
-
-
-def get_email_count(conn, folder_name):
-    """Return the IMAP message count for a folder, or None on error."""
-    try:
-        # Select folder in read-only mode
-        # Quote folder name handles spaces
-        typ, data = conn.select(f'"{folder_name}"', readonly=True)
-        if typ != "OK":
-            return None
-
-        # SELECT command returns the number of messages in data[0]
-        # data[0] is bytes, e.g. b'123'
-        if data and data[0]:
-            return int(data[0])
-        return 0
-
-    except Exception:
-        # print(f"Error checking {folder_name}: {e}")
-        return None
-
-
-def main():
-    default_src_path = os.getenv("SRC_LOCAL_PATH")
-    default_dest_path = os.getenv("DEST_LOCAL_PATH")
-
-    # Phase 1: determine whether each side is local
-    pre_parser = argparse.ArgumentParser(add_help=False)
-    pre_parser.add_argument("--src-path", default=default_src_path)
-    pre_parser.add_argument("--dest-path", default=default_dest_path)
-    pre_args, _ = pre_parser.parse_known_args()
-
-    src_requires_imap = not bool(pre_args.src_path)
-    dest_requires_imap = not bool(pre_args.dest_path)
-
-    # Phase 2: full parser with conditional requirements
-    parser = argparse.ArgumentParser(description="Compare email counts between two IMAP accounts.")
-    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
-    parser.add_argument(
-        "--src-path",
-        default=default_src_path,
-        help="Source local folder (backup root). If set, IMAP source args are ignored.",
-    )
-    parser.add_argument(
-        "--dest-path",
-        default=default_dest_path,
-        help="Destination local folder (backup root). If set, IMAP destination args are ignored.",
-    )
-
-    # Source args
-    default_src_host = os.getenv("SRC_IMAP_HOST")
-    default_src_user = os.getenv("SRC_IMAP_USERNAME")
-    default_src_pass = os.getenv("SRC_IMAP_PASSWORD")
-    default_src_client_id = os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--src-host",
-        default=default_src_host,
-        required=src_requires_imap and not bool(default_src_host),
-        help="Source IMAP Server (or SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--src-user",
-        default=default_src_user,
-        required=src_requires_imap and not bool(default_src_user),
-        help="Source Username (or SRC_IMAP_USERNAME)",
-    )
-    src_auth_required = src_requires_imap and not bool(default_src_pass or default_src_client_id)
-    src_auth = parser.add_mutually_exclusive_group(required=src_auth_required)
-    src_auth.add_argument("--src-pass", default=default_src_pass, help="Source Password (or SRC_IMAP_PASSWORD)")
-    src_auth.add_argument(
-        "--src-oauth2-client-id",
-        default=default_src_client_id,
-        dest="src_client_id",
-        help="Source OAuth2 Client ID (or SRC_OAUTH2_CLIENT_ID)",
-    )
-    parser.add_argument(
-        "--src-oauth2-client-secret",
-        default=os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="src_client_secret",
-        help="Source OAuth2 Client Secret (if required) (or SRC_OAUTH2_CLIENT_SECRET)",
-    )
-
-    # Dest args
-    default_dest_host = os.getenv("DEST_IMAP_HOST")
-    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
-    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
-    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--dest-host",
-        default=default_dest_host,
-        required=dest_requires_imap and not bool(default_dest_host),
-        help="Destination IMAP Server (or DEST_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--dest-user",
-        default=default_dest_user,
-        required=dest_requires_imap and not bool(default_dest_user),
-        help="Destination Username (or DEST_IMAP_USERNAME)",
-    )
-    dest_auth_required = dest_requires_imap and not bool(default_dest_pass or default_dest_client_id)
-    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
-    dest_auth.add_argument(
-        "--dest-pass",
-        default=default_dest_pass,
-        help="Destination Password (or DEST_IMAP_PASSWORD)",
-    )
-    dest_auth.add_argument(
-        "--dest-oauth2-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
-    )
-    dest_auth.add_argument(
-        "--dest-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help=argparse.SUPPRESS,
-    )
-    parser.add_argument(
-        "--dest-oauth2-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
-    )
-    parser.add_argument(
-        "--dest-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help=argparse.SUPPRESS,
-    )
-
-    args = parser.parse_args()
-
-    src_is_local = bool(args.src_path)
-    dest_is_local = bool(args.dest_path)
-
-    SRC_HOST = args.src_host
-    SRC_USER = args.src_user
-    DEST_HOST = args.dest_host
-    DEST_USER = args.dest_user
-
-    # Acquire OAuth2 tokens if configured
-    src_oauth2_token = None
-    src_oauth2_provider = None
-    if not src_is_local and args.src_client_id:
-        src_oauth2_token, src_oauth2_provider = imap_oauth2.acquire_token(
-            SRC_HOST, args.src_client_id, SRC_USER, args.src_client_secret, "source"
-        )
-
-    dest_oauth2_token = None
-    dest_oauth2_provider = None
-    if not dest_is_local and args.dest_client_id:
-        dest_oauth2_token, dest_oauth2_provider = imap_oauth2.acquire_token(
-            DEST_HOST, args.dest_client_id, DEST_USER, args.dest_client_secret, "destination"
-        )
-
-    print("\n--- Configuration Summary ---")
-    if src_is_local:
-        print(f"Source (Local)  : {args.src_path}")
-    else:
-        print(f"Source Host     : {args.src_host}")
-        print(f"Source User     : {args.src_user}")
-        print(f"Source Auth     : {imap_oauth2.auth_description(src_oauth2_provider)}")
-
-    if dest_is_local:
-        print(f"Destination (Local): {args.dest_path}")
-    else:
-        print(f"Destination Host: {args.dest_host}")
-        print(f"Destination User: {args.dest_user}")
-        print(f"Destination Auth: {imap_oauth2.auth_description(dest_oauth2_provider)}")
-    print("-----------------------------\n")
-
-    src = None
-    dest = None
-
-    try:
-        if not src_is_local:
-            # Connect to Source
-            print("Connecting to Source...")
-            src = imap_common.get_imap_connection(args.src_host, args.src_user, args.src_pass, src_oauth2_token)
-            if not src:
-                return
-
-        if not dest_is_local:
-            # Connect to Dest
-            print("Connecting to Destination...")
-            dest = imap_common.get_imap_connection(args.dest_host, args.dest_user, args.dest_pass, dest_oauth2_token)
-            if not dest:
-                return
-
-        # List Source Folders
-        print("Listing folders in Source...")
-        if src_is_local:
-            folders = imap_common.list_local_folders(args.src_path)
-        else:
-            folders = imap_common.list_selectable_folders(src)
-
-        if not folders:
-            print("Failed to list source folders.")
-            return
-
-        # Prepare Table Header
-        header = f"{'Folder Name':<40} | {'Source':>10} | {'Dest':>10} | {'Diff':>10}"
-        print("-" * len(header))
-        print(header)
-        print("-" * len(header))
-
-        total_src = 0
-        total_dest = 0
-
-        # Iterate through Source folders
-        for folder_name in folders:
-            # Get Counts
-            if src_is_local:
-                src_count = imap_common.get_local_email_count(args.src_path, folder_name)
-            else:
-                src_count = get_email_count(src, folder_name)
-
-            if dest_is_local:
-                dest_count = imap_common.get_local_email_count(args.dest_path, folder_name)
-            else:
-                dest_count = get_email_count(dest, folder_name)
-
-            # Format for display
-            src_str = str(src_count) if src_count is not None else "Err"
-            dest_str = str(dest_count) if dest_count is not None else "N/A"  # N/A usually means folder doesn't exist
-
-            diff_str = ""
-            if src_count is not None and dest_count is not None:
-                diff = src_count - dest_count
-                diff_str = str(diff)
-                total_src += src_count
-                total_dest += dest_count
-            elif src_count is not None:
-                total_src += src_count
-
-            print(f"{folder_name:<40} | {src_str:>10} | {dest_str:>10} | {diff_str:>10}")
-
-        print("-" * len(header))
-        print(f"{'TOTAL':<40} | {total_src:>10} | {total_dest:>10} | {total_src - total_dest:>10}")
-
-    except KeyboardInterrupt:
-        # Re-raise to be handled by the outer block, but ensure finally runs
-        raise
-
-    finally:
-        # Check source connection state and logout if possible
-        if src:
-            try:
-                src.logout()
-            except BaseException:
-                pass
-
-        # Check dest connection state and logout if possible
-        if dest:
-            try:
-                dest.logout()
-            except BaseException:
-                pass
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        sys.exit(1)

src/imap_count.py

@@ -1,247 +0,0 @@
-"""IMAP Email Counting Script.
-
-Counts emails per folder from either:
-- An IMAP account, or
-- A local backup folder created by ``imap_backup.py`` (counts ``.eml`` files).
-
-Configuration (Environment Variables):
-    IMAP_HOST        : IMAP Host (e.g., imap.gmail.com)
-    IMAP_USERNAME    : Username/Email
-    IMAP_PASSWORD    : Password (or App Password)
-
-    OAuth2 (Optional - instead of password):
-    OAUTH2_CLIENT_ID        : OAuth2 Client ID
-    OAUTH2_CLIENT_SECRET    : OAuth2 Client Secret (required for Google)
-    SRC_OAUTH2_CLIENT_ID    : Alternate OAuth2 client ID env var
-    SRC_OAUTH2_CLIENT_SECRET: Alternate OAuth2 client secret env var
-
-Local backup counting:
-    BACKUP_LOCAL_PATH : Local backup root (preferred)
-    SRC_LOCAL_PATH    : Alternate local backup root
-
-Examples:
-    # Count an IMAP account
-    export IMAP_HOST="imap.gmail.com"
-    export IMAP_USERNAME="user@gmail.com"
-    export IMAP_PASSWORD="secretpassword"
-    python3 imap_count.py
-
-    # Count an IMAP account using OAuth2
-    export IMAP_HOST="imap.gmail.com"
-    export IMAP_USERNAME="user@gmail.com"
-    export OAUTH2_CLIENT_ID="your-client-id"
-    export OAUTH2_CLIENT_SECRET="your-client-secret"  # Required for Google
-    python3 imap_count.py
-
-    # Count a local backup
-    python3 imap_count.py --path "./my_backup"
-
-    # Or set a default local backup path via env var
-    export BACKUP_LOCAL_PATH="./my_backup"
-    python3 imap_count.py
-"""
-
-import argparse
-import imaplib
-import os
-import sys
-from typing import Optional
-
-from auth import imap_oauth2
-from utils import imap_common
-
-
-def count_emails(imap_server, username, password=None, oauth2_token=None):
-    try:
-        # Connect to the IMAP server (using SSL)
-        print(f"Connecting to {imap_server}...")
-        mail = imap_common.get_imap_connection(imap_server, username, password, oauth2_token)
-        if not mail:
-            return
-
-        # List all mailboxes
-        print("Listing mailboxes...")
-        folders = imap_common.list_selectable_folders(mail)
-
-        if not folders:
-            print("Failed to list mailboxes.")
-            return
-
-        total_all_folders = 0
-        print(f"{'Folder Name':<40} {'Count':>10}")
-        print("-" * 52)
-
-        for folder_name in folders:
-            display_name = folder_name
-
-            try:
-                # Select the mailbox (read-only is sufficient for counting)
-                # folder_name extracted from list usually handles quotes correctly for select
-                rv, _ = mail.select(f'"{folder_name}"', readonly=True)
-                if rv != "OK":
-                    print(f"{display_name:<40} {'Skipped':>10}")
-                    continue
-
-                # Search for all emails
-                status, data = mail.search(None, "ALL")
-
-                if status == "OK":
-                    # data[0] is space separated IDs
-                    email_ids = data[0].split()
-                    count = len(email_ids)
-                    print(f"{display_name:<40} {count:>10}")
-                    total_all_folders += count
-                else:
-                    print(f"{display_name:<40} {'Error':>10}")
-
-            except imaplib.IMAP4.error:
-                print(f"{display_name:<40} {'Error':>10}")
-
-        print("-" * 52)
-        print(f"{'TOTAL':<40} {total_all_folders:>10}")
-
-        # Logout
-        mail.logout()
-
-    except imaplib.IMAP4.error as e:
-        print(f"IMAP Error: {e}")
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-
-def count_local_emails(local_path: str) -> None:
-    print(f"Scanning local backup: {local_path}")
-
-    folders = imap_common.list_local_folders(local_path)
-    if not folders:
-        print("No folders found.")
-        return
-
-    total_all_folders = 0
-    print(f"{'Folder Name':<40} {'Count':>10}")
-    print("-" * 52)
-
-    for folder_name in folders:
-        count = imap_common.get_local_email_count(local_path, folder_name)
-        if count is None:
-            print(f"{folder_name:<40} {'N/A':>10}")
-            continue
-
-        print(f"{folder_name:<40} {count:>10}")
-        total_all_folders += count
-
-    print("-" * 52)
-    print(f"{'TOTAL':<40} {total_all_folders:>10}")
-
-
-def main(argv: Optional[list[str]] = None) -> None:
-    # Phase 1: determine whether we're in local mode (--path)
-    default_path = os.getenv("BACKUP_LOCAL_PATH") or os.getenv("SRC_LOCAL_PATH")
-    pre_parser = argparse.ArgumentParser(add_help=False)
-    pre_parser.add_argument("--path", default=default_path)
-    pre_args, _ = pre_parser.parse_known_args(argv)
-    require_imap = not bool(pre_args.path)
-
-    # Phase 2: full parser with conditional requirements
-    parser = argparse.ArgumentParser(description="Count emails in IMAP account.")
-    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
-
-    parser.add_argument(
-        "--path",
-        default=default_path,
-        help="Local backup root to count (counts .eml files per folder). If set, IMAP args are ignored.",
-    )
-
-    # Try to unify var names for defaults. Priority: IMAP_* > SRC_IMAP_* > None
-    default_host = os.getenv("IMAP_HOST") or os.getenv("SRC_IMAP_HOST")
-    default_user = os.getenv("IMAP_USERNAME") or os.getenv("SRC_IMAP_USERNAME")
-    default_pass = os.getenv("IMAP_PASSWORD") or os.getenv("SRC_IMAP_PASSWORD")
-    default_client_id = os.getenv("OAUTH2_CLIENT_ID") or os.getenv("SRC_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--host",
-        default=default_host,
-        required=require_imap and not bool(default_host),
-        help="IMAP Server (or IMAP_HOST / SRC_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--user",
-        default=default_user,
-        required=require_imap and not bool(default_user),
-        help="Username (or IMAP_USERNAME / SRC_IMAP_USERNAME)",
-    )
-
-    auth_required = require_imap and not bool(default_pass or default_client_id)
-    auth_group = parser.add_mutually_exclusive_group(required=auth_required)
-    auth_group.add_argument(
-        "--pass", dest="password", default=default_pass, help="Password (or IMAP_PASSWORD / SRC_IMAP_PASSWORD)"
-    )
-    auth_group.add_argument(
-        "--oauth2-client-id",
-        default=default_client_id,
-        dest="client_id",
-        help="OAuth2 Client ID (or OAUTH2_CLIENT_ID / SRC_OAUTH2_CLIENT_ID)",
-    )
-    auth_group.add_argument(
-        "--client-id",
-        default=default_client_id,
-        dest="client_id",
-        help=argparse.SUPPRESS,
-    )
-
-    parser.add_argument(
-        "--oauth2-client-secret",
-        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="client_secret",
-        help="OAuth2 Client Secret (if required) (or OAUTH2_CLIENT_SECRET / SRC_OAUTH2_CLIENT_SECRET)",
-    )
-    parser.add_argument(
-        "--client-secret",
-        default=os.getenv("OAUTH2_CLIENT_SECRET") or os.getenv("SRC_OAUTH2_CLIENT_SECRET"),
-        dest="client_secret",
-        help=argparse.SUPPRESS,
-    )
-
-    args = parser.parse_args(argv)
-
-    if args.path:
-        if not os.path.isdir(args.path):
-            print(f"Error: Local path does not exist or is not a directory: {args.path}")
-            sys.exit(1)
-
-        print("\n--- Configuration Summary ---")
-        print(f"Local Path      : {args.path}")
-        print("-----------------------------\n")
-        count_local_emails(args.path)
-        raise SystemExit(0)
-
-    IMAP_SERVER = args.host
-    USERNAME = args.user
-    PASSWORD = args.password
-
-    # Acquire OAuth2 token if configured
-    oauth2_token = None
-    oauth2_provider = None
-    if args.client_id:
-        oauth2_token, oauth2_provider = imap_oauth2.acquire_token(
-            IMAP_SERVER, args.client_id, USERNAME, args.client_secret
-        )
-
-    print("\n--- Configuration Summary ---")
-    print(f"Host            : {IMAP_SERVER}")
-    print(f"User            : {USERNAME}")
-    print(f"Auth Method     : {imap_oauth2.auth_description(oauth2_provider)}")
-    print("-----------------------------\n")
-
-    count_emails(IMAP_SERVER, USERNAME, PASSWORD, oauth2_token)
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        sys.exit(1)

src/imap_oauth2.py

@@ -10,7 +10,8 @@ Provider is auto-detected from the IMAP host string.
 import sys
 import threading
 
-from auth import oauth2_google, oauth2_microsoft
+import oauth2_google
+import oauth2_microsoft
 
 # Re-export caches for test access
 _msal_app_cache = oauth2_microsoft._msal_app_cache

src/imap_restore.py

@@ -1,976 +0,0 @@
-"""
-IMAP Email Restore Script
-
-Restores emails from a local backup to an IMAP account.
-Reads .eml files from a local directory and uploads them to the destination server.
-
-Features:
-- Folder Restoration: Recreates the folder structure from the backup.
-- Gmail Labels Restoration: Uses labels_manifest.json to apply Gmail labels.
-- Incremental Restore: Skips emails that already exist (based on Message-ID).
-- Parallel Processing: Uses multithreading for fast uploads.
-- Date Preservation: Restores emails with their original dates.
-
-Configuration (Environment Variables):
-    DEST_IMAP_HOST, DEST_IMAP_USERNAME: Destination credentials.
-    DEST_IMAP_PASSWORD: Destination password (or App Password).
-
-    OAuth2 (Optional - instead of password):
-    DEST_OAUTH2_CLIENT_ID: OAuth2 Client ID
-    DEST_OAUTH2_CLIENT_SECRET: OAuth2 Client Secret (required for Google)
-
-  BACKUP_LOCAL_PATH: Source local directory containing the backup.
-  MAX_WORKERS: Number of concurrent threads (default: 4).
-  BATCH_SIZE: Number of emails to process per batch (default: 10).
-  APPLY_LABELS: Set to "true" to apply Gmail labels from manifest. Default is "false".
-  APPLY_FLAGS: Set to "true" to apply IMAP flags from manifest. Default is "false".
-  GMAIL_MODE: Set to "true" for Gmail restore mode. Default is "false".
-  DEST_DELETE: Set to "true" to delete emails from destination not found in local backup.
-              Default is "false".
-
-Usage:
-    python3 imap_restore.py \
-        --src-path "./my_backup" \
-        --dest-host "imap.gmail.com" \
-        --dest-user "you@gmail.com" \
-        --dest-pass "your-app-password"
-
-Gmail Labels Restoration:
-    python3 imap_restore.py \
-        --src-path "./gmail_backup" \
-        --dest-host "imap.gmail.com" \
-        --dest-user "you@gmail.com" \
-        --dest-pass "your-app-password" \
-        --apply-labels
-  This uploads emails and applies labels from labels_manifest.json to recreate
-  the original Gmail label structure.
-"""
-
-import argparse
-import concurrent.futures
-import os
-import sys
-import threading
-import time
-from email import policy
-from email.parser import BytesParser
-from email.utils import parsedate_to_datetime
-from enum import Enum
-from typing import Optional
-
-from auth import imap_oauth2
-from core import imap_session
-from providers import provider_gmail
-from utils import imap_common, restore_cache
-
-
-class UploadResult(Enum):
-    """Result of an email upload operation."""
-
-    SUCCESS = "success"
-    ALREADY_EXISTS = "already_exists"
-    FAILURE = "failure"
-
-
-# Defaults
-MAX_WORKERS = 4  # Lower default for restore to avoid rate limits
-BATCH_SIZE = 10
-
-# Thread-local storage
-thread_local = threading.local()
-safe_print = imap_common.safe_print
-
-
-def get_flags_from_manifest(manifest, message_id):
-    """
-    Extract IMAP flags string from manifest entry.
-    Returns flags string like "\\Seen \\Flagged" or None.
-    """
-    if not message_id or message_id not in manifest:
-        return None
-
-    entry = manifest[message_id]
-
-    if isinstance(entry, dict) and "flags" in entry:
-        flags = entry.get("flags", [])
-        if flags:
-            return " ".join(flags)
-
-    return None
-
-
-def parse_eml_file(file_path):
-    """
-    Parse an .eml file and extract metadata.
-    Returns (message_id, date_str, raw_content, subject) or (None, None, None, None) on error.
-    """
-    try:
-        with open(file_path, "rb") as f:
-            raw_content = f.read()
-
-        # Use compat32 to preserve raw headers with continuation lines
-        parser = BytesParser(policy=policy.compat32)
-        msg = parser.parsebytes(raw_content, headersonly=True)
-
-        message_id = imap_common.decode_message_id(msg.get("Message-ID"))
-        raw_subject = msg.get("Subject")
-        subject = imap_common.decode_mime_header(raw_subject) if raw_subject else "(No Subject)"
-        date_header = msg.get("Date")
-
-        # Parse date for IMAP INTERNALDATE
-        date_str = None
-        if date_header:
-            try:
-                dt = parsedate_to_datetime(date_header)
-                # Format: "DD-Mon-YYYY HH:MM:SS +ZZZZ"
-                date_str = dt.strftime('"%d-%b-%Y %H:%M:%S %z"')
-            except Exception:
-                pass
-
-        return message_id, date_str, raw_content, subject
-
-    except Exception as e:
-        safe_print(f"Error parsing {file_path}: {e}")
-        return None, None, None, None
-
-
-def get_eml_files(folder_path):
-    """
-    Get all .eml files in a folder.
-    Returns list of (file_path, filename) tuples.
-    """
-    eml_files = []
-    if not os.path.exists(folder_path):
-        return eml_files
-
-    try:
-        for filename in os.listdir(folder_path):
-            if filename.endswith(".eml"):
-                file_path = os.path.join(folder_path, filename)
-                eml_files.append((file_path, filename))
-    except Exception as e:
-        safe_print(f"Error listing {folder_path}: {e}")
-
-    return eml_files
-
-
-def upload_email(dest, folder_name, raw_content, date_str, flags=None):
-    """
-    Upload a single email to the destination folder.
-    Returns UploadResult enum: SUCCESS or FAILURE.
-
-    Callers are expected to check for duplicates via load_folder_msg_ids()
-    before calling this function, and to ensure the folder exists.
-
-    Args:
-        flags: Optional string of IMAP flags like "\\Seen" for read emails.
-    """
-    try:
-        success = imap_common.append_email(
-            dest,
-            folder_name,
-            raw_content,
-            date_str,
-            flags,
-            ensure_folder=False,
-        )
-        return UploadResult.SUCCESS if success else UploadResult.FAILURE
-
-    except Exception as e:
-        safe_print(f"Error uploading to {folder_name}: {e}")
-        return UploadResult.FAILURE
-
-
-def get_labels_from_manifest(manifest, message_id):
-    """
-    Get the list of labels for a message from the manifest.
-    Returns list of label strings or empty list.
-    """
-    if not message_id or message_id not in manifest:
-        return []
-
-    entry = manifest[message_id]
-    if isinstance(entry, dict):
-        return entry.get("labels", [])
-    elif isinstance(entry, list):
-        return entry  # Old format: list of labels
-    return []
-
-
-def process_restore_batch(
-    eml_files,
-    folder_name,
-    dest_conf,
-    manifest,
-    apply_labels,
-    apply_flags,
-    full_restore: bool = False,
-    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = None,
-    existing_dest_msg_ids_lock: Optional[threading.Lock] = None,
-    progress_cache_data: Optional[dict] = None,
-    progress_cache_lock: Optional[threading.Lock] = None,
-    progress_cache_path: Optional[str] = None,
-    dest_host: Optional[str] = None,
-    dest_user: Optional[str] = None,
-):
-    """
-    Process a batch of .eml files for restoration.
-
-    Args:
-        folder_name: Target folder, or "__GMAIL_MODE__" for per-email folder selection
-        manifest: Combined manifest with labels and/or flags
-        apply_labels: Whether to apply Gmail labels from manifest
-        apply_flags: Whether to apply IMAP flags from manifest
-    """
-    dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
-    if not dest:
-        safe_print("Error: Could not establish connection for batch.")
-        return
-
-    gmail_mode = folder_name == "__GMAIL_MODE__"
-
-    for file_path, filename in eml_files:
-        # Proactively refresh token if needed
-        dest = imap_session.get_thread_connection(thread_local, "dest", dest_conf)
-        if not dest:
-            safe_print(f"ERROR: Connection lost for {filename}")
-            return
-
-        try:
-            message_id, date_str, raw_content, subject = parse_eml_file(file_path)
-            if raw_content is None:
-                continue  # No content, skip to next file
-
-            size = len(raw_content)
-            size_str = f"{size / 1024:.1f}KB"
-
-            # Truncate subject for display
-            display_subject = (subject[:40] + "...") if len(subject) > 40 else subject
-
-            # Get flags from manifest if apply_flags is enabled
-            flags = None
-            if apply_flags:
-                flags = get_flags_from_manifest(manifest, message_id)
-
-            # Get labels for this message
-            labels = get_labels_from_manifest(manifest, message_id) if apply_labels else []
-
-            # Determine target folder and remaining labels
-            if gmail_mode:
-                target_folder, remaining_labels = provider_gmail.resolve_target(labels)
-            else:
-                target_folder = folder_name
-                remaining_labels = labels
-
-            existing_dest_msg_ids = imap_common.load_folder_msg_ids(
-                dest,
-                target_folder,
-                existing_dest_msg_ids_by_folder,
-                existing_dest_msg_ids_lock,
-                progress_cache_data,
-                progress_cache_lock,
-                dest_host,
-                dest_user,
-            )
-
-            # Check if email already exists on destination using pre-fetched set
-            email_already_on_dest = (
-                message_id and existing_dest_msg_ids is not None and message_id in existing_dest_msg_ids
-            )
-
-            # Incremental default: skip entirely if already present
-            if email_already_on_dest and not full_restore:
-                safe_print(f"[{target_folder}] SKIP (already present) | {size_str:<8} | {display_subject}")
-                continue  # Skip to next file
-
-            if email_already_on_dest:
-                # Full restore: treat as existing for flag sync
-                upload_result = UploadResult.ALREADY_EXISTS
-            else:
-                # Upload — skip per-message SEARCH since we have a pre-fetched set
-                upload_result = upload_email(
-                    dest,
-                    target_folder,
-                    raw_content,
-                    date_str,
-                    flags,
-                )
-
-            # Only record progress when upload succeeds or email already exists.
-            # Failed uploads should not be marked as processed to allow retry on next run.
-            if upload_result in (UploadResult.SUCCESS, UploadResult.ALREADY_EXISTS):
-                restore_cache.record_progress(
-                    message_id=message_id,
-                    folder_name=target_folder,
-                    existing_dest_msg_ids=existing_dest_msg_ids,
-                    existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
-                    progress_cache_path=progress_cache_path,
-                    progress_cache_data=progress_cache_data,
-                    progress_cache_lock=progress_cache_lock,
-                    dest_host=dest_host,
-                    dest_user=dest_user,
-                    log_fn=safe_print,
-                )
-
-            if upload_result == UploadResult.ALREADY_EXISTS:
-                safe_print(f"[{target_folder}] SKIP (exists) | {size_str:<8} | {display_subject}")
-                # Full restore preserves legacy behavior: sync flags on existing email if requested
-                if full_restore and apply_flags and flags and message_id:
-                    imap_common.sync_flags_on_existing(dest, target_folder, message_id, flags, size)
-            elif upload_result == UploadResult.SUCCESS:
-                safe_print(f"[{target_folder}] UPLOADED      | {size_str:<8} | {display_subject}")
-                # Show applied flags in same style as labels
-                if flags:
-                    for flag in flags.split():
-                        safe_print(f"  -> Applied flag: {flag}")
-            else:  # FAILURE
-                safe_print(f"[{target_folder}] FAILED        | {size_str:<8} | {display_subject}")
-
-            # Apply remaining Gmail labels:
-            # - Full restore: apply/sync labels even for existing emails
-            # - Incremental (default): apply labels only for newly uploaded emails
-            if apply_labels and remaining_labels and (upload_result == UploadResult.SUCCESS or full_restore):
-                for label in remaining_labels:
-                    label_folder = provider_gmail.label_to_folder(label)
-
-                    # Skip if this is the same as the target folder
-                    if label_folder == target_folder:
-                        continue
-
-                    # Skip system folders we can't upload to
-                    if label_folder in (
-                        provider_gmail.GMAIL_ALL_MAIL,
-                        provider_gmail.GMAIL_SPAM,
-                        provider_gmail.GMAIL_TRASH,
-                    ):
-                        continue
-
-                    try:
-                        # Get or fetch Message-IDs for label folder (one-time server fetch per folder)
-                        label_folder_msg_ids = imap_common.load_folder_msg_ids(
-                            dest,
-                            label_folder,
-                            existing_dest_msg_ids_by_folder,
-                            existing_dest_msg_ids_lock,
-                            progress_cache_data,
-                            progress_cache_lock,
-                            dest_host,
-                            dest_user,
-                        )
-
-                        # Check duplicate using pre-fetched set instead of per-message SEARCH
-                        label_already_exists = (
-                            label_folder_msg_ids is not None and message_id and message_id in label_folder_msg_ids
-                        )
-
-                        if not label_already_exists:
-                            append_success = imap_common.append_email(
-                                dest,
-                                label_folder,
-                                raw_content,
-                                date_str,
-                                flags,
-                                ensure_folder=False,
-                            )
-                            if append_success:
-                                restore_cache.record_progress(
-                                    message_id=message_id,
-                                    folder_name=label_folder,
-                                    existing_dest_msg_ids=label_folder_msg_ids,
-                                    existing_dest_msg_ids_lock=existing_dest_msg_ids_lock,
-                                    progress_cache_path=progress_cache_path,
-                                    progress_cache_data=progress_cache_data,
-                                    progress_cache_lock=progress_cache_lock,
-                                    dest_host=dest_host,
-                                    dest_user=dest_user,
-                                    log_fn=safe_print,
-                                )
-                                safe_print(f"  -> Applied label: {label}")
-                            else:
-                                safe_print(f"  -> Failed to apply label {label} (will retry on next restore)")
-                        # If email exists in this label folder, sync flags (full restore only)
-                        elif full_restore and apply_flags and flags:
-                            imap_common.sync_flags_on_existing(dest, label_folder, message_id, flags, size)
-                    except Exception as e:
-                        if imap_oauth2.is_auth_error(e):
-                            raise
-                        safe_print(f"  -> Error applying label {label}: {e}")
-
-        except Exception as e:
-            safe_print(f"Error processing {filename}: {e}")
-
-
-def get_local_message_ids(local_folder_path):
-    """
-    Get a set of Message-IDs from all .eml files in a local folder.
-    Used for destination deletion sync.
-    """
-    message_ids = set()
-    eml_files = get_eml_files(local_folder_path)
-    for file_path, _filename in eml_files:
-        message_id, _, _, _ = parse_eml_file(file_path)
-        if message_id:
-            message_ids.add(message_id)
-    return message_ids
-
-
-def pre_filter_eml_files(eml_files, dest_msg_ids):
-    """Filter out .eml files whose Message-IDs already exist in the destination."""
-    safe_print("Pre-filtering duplicates...")
-    files_to_restore = []
-    skipped = 0
-    for file_path, filename in eml_files:
-        msg_id = imap_common.extract_message_id_from_eml(file_path)
-        if msg_id and msg_id in dest_msg_ids:
-            skipped += 1
-        else:
-            files_to_restore.append((file_path, filename))
-    safe_print(f"Skipping {skipped} duplicates, {len(files_to_restore)} to restore.")
-    return files_to_restore
-
-
-def restore_folder(
-    folder_name,
-    local_folder_path,
-    dest_conf,
-    manifest,
-    apply_labels,
-    apply_flags,
-    dest_delete=False,
-    full_restore: bool = False,
-    cache_root: Optional[str] = None,
-    progress_cache_file: Optional[str] = None,
-    progress_cache_data: Optional[dict] = None,
-    progress_cache_lock: Optional[threading.Lock] = None,
-):
-    """
-    Restore all emails from a local folder to the destination IMAP server.
-    """
-    safe_print(f"--- Restoring Folder: {folder_name} ---")
-
-    eml_files = get_eml_files(local_folder_path)
-    if not eml_files:
-        safe_print(f"No .eml files found in {folder_name}")
-        # Even if empty, check for orphans to delete
-        if dest_delete:
-            dest = imap_common.get_imap_connection_from_conf(dest_conf)
-            if dest:
-                imap_common.delete_orphan_emails(dest, folder_name, set())
-                dest.logout()
-        return
-
-    safe_print(f"Found {len(eml_files)} emails to restore.")
-
-    cache_root = cache_root or local_folder_path
-    cache_path = progress_cache_file
-    cache_data = progress_cache_data
-    cache_lock = progress_cache_lock
-    if cache_path is None or cache_data is None or cache_lock is None:
-        cache_path, cache_data, cache_lock = imap_common.load_progress_cache(
-            cache_root,
-            dest_conf["host"],
-            dest_conf["user"],
-            log_fn=safe_print,
-        )
-
-    # Incremental mode uses cached Message-IDs to skip already-processed emails.
-    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {folder_name: set()}
-    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
-    try:
-        existing_dest_msg_ids_by_folder[folder_name] = restore_cache.get_cached_message_ids(
-            cache_data,
-            cache_lock,
-            dest_conf["host"],
-            dest_conf["user"],
-            folder_name,
-        )
-        safe_print(f"Cache has {len(existing_dest_msg_ids_by_folder[folder_name])} Message-IDs for this folder.")
-    except Exception as e:
-        # Fall back to an empty cache for this folder if reading cached Message-IDs fails.
-        safe_print(f"Warning: Failed to load cached Message-IDs for folder '{folder_name}': {e}")
-        existing_dest_msg_ids_by_folder[folder_name] = set()
-
-    # If dest_delete enabled, get local Message-IDs for comparison
-    local_msg_ids = None
-    if dest_delete:
-        safe_print("Building local Message-ID index for sync...")
-        local_msg_ids = get_local_message_ids(local_folder_path)
-        safe_print(f"Found {len(local_msg_ids)} unique Message-IDs in local backup.")
-
-    # Pre-fetch destination Message-IDs and filter duplicates (non-Gmail mode only)
-    gmail_mode = folder_name == "__GMAIL_MODE__"
-    files_to_restore = eml_files
-
-    if not gmail_mode:
-        dest_msg_ids = set()
-        try:
-            dest_tmp = imap_common.get_imap_connection_from_conf(dest_conf)
-            if dest_tmp:
-                # Ensure folder exists before selecting
-                if folder_name.upper() != "INBOX":
-                    try:
-                        dest_tmp.create(f'"{folder_name}"')
-                    except Exception:
-                        pass
-                dest_tmp.select(f'"{folder_name}"')
-                dest_msg_ids = set(imap_common.get_message_ids_in_folder(dest_tmp).values())
-                dest_tmp.logout()
-        except Exception:
-            dest_msg_ids = set()
-
-        safe_print(f"{len(dest_msg_ids)} existing messages in destination.")
-
-        # Update destination Message-IDs with what we processed
-        if not full_restore:
-            with existing_dest_msg_ids_lock:
-                existing_dest_msg_ids_by_folder.setdefault(folder_name, set()).update(dest_msg_ids)
-        else:
-            existing_dest_msg_ids_by_folder[folder_name] = dest_msg_ids
-        dest_msg_ids = existing_dest_msg_ids_by_folder[folder_name]
-
-        # Pre-filter files to skip duplicates (skip in full_restore: need to process all for flag/label sync)
-        if dest_msg_ids and not full_restore:
-            files_to_restore = pre_filter_eml_files(eml_files, dest_msg_ids)
-
-    if not files_to_restore:
-        safe_print("No new emails to restore.")
-        if dest_delete and local_msg_ids is not None:
-            safe_print("Syncing destination: removing emails not in local backup...")
-            dest = imap_session.ensure_connection(None, dest_conf)
-            if dest:
-                imap_common.delete_orphan_emails(dest, folder_name, local_msg_ids)
-                dest.logout()
-
-        restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
-        return
-
-    safe_print(f"Starting parallel restore of {len(files_to_restore)} emails...")
-
-    # Create batches
-    batches = [files_to_restore[i : i + BATCH_SIZE] for i in range(0, len(files_to_restore), BATCH_SIZE)]
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        futures = []
-        for batch in batches:
-            futures.append(
-                executor.submit(
-                    process_restore_batch,
-                    batch,
-                    folder_name,
-                    dest_conf,
-                    manifest,
-                    apply_labels,
-                    apply_flags,
-                    full_restore,
-                    existing_dest_msg_ids_by_folder,
-                    existing_dest_msg_ids_lock,
-                    cache_data,
-                    cache_lock,
-                    cache_path,
-                    dest_conf.get("host"),
-                    dest_conf.get("user"),
-                )
-            )
-
-        for future in concurrent.futures.as_completed(futures):
-            try:
-                future.result()
-            except Exception as e:
-                safe_print(f"Batch error: {e}")
-
-    # Delete orphan emails from destination if enabled
-    if dest_delete and local_msg_ids is not None:
-        safe_print("Syncing destination: removing emails not in local backup...")
-        dest = imap_session.ensure_connection(None, dest_conf)
-        if dest:
-            imap_common.delete_orphan_emails(dest, folder_name, local_msg_ids)
-            dest.logout()
-
-    # Force-flush progress cache at end.
-    restore_cache.maybe_save_dest_index_cache(cache_path, cache_data, cache_lock, force=True)
-
-
-def restore_gmail_with_labels(
-    local_path,
-    dest_conf,
-    manifest,
-    apply_flags,
-    full_restore: bool = False,
-    progress_cache_file: Optional[str] = None,
-    progress_cache_data: Optional[dict] = None,
-    progress_cache_lock: Optional[threading.Lock] = None,
-):
-    """
-    Special restoration mode for Gmail: Upload emails to their first label folder
-    and then apply additional labels from the manifest.
-
-    This avoids putting all emails in INBOX - emails only appear in INBOX
-    if they originally had the INBOX label.
-    """
-    # Find the All Mail folder in the backup
-    all_mail_path = os.path.join(local_path, "[Gmail]", "All Mail")
-    if not os.path.exists(all_mail_path):
-        # Try without subfolder structure
-        all_mail_path = local_path
-
-    safe_print("--- Gmail Restore with Labels ---")
-    safe_print(f"Source path: {all_mail_path}")
-    safe_print(f"Entries in manifest: {len(manifest)}")
-
-    eml_files = get_eml_files(all_mail_path)
-    if not eml_files:
-        safe_print("No .eml files found for restoration.")
-        return
-
-    safe_print(f"Found {len(eml_files)} emails to restore.\n")
-
-    # Process in batches
-    total = len(eml_files)
-    start_time = time.time()
-
-    batches = [eml_files[i : i + BATCH_SIZE] for i in range(0, len(eml_files), BATCH_SIZE)]
-
-    cache_path = progress_cache_file
-    if cache_path is None or progress_cache_data is None or progress_cache_lock is None:
-        cache_path, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
-            local_path,
-            dest_conf["host"],
-            dest_conf["user"],
-            log_fn=safe_print,
-        )
-    safe_print(
-        "Cache will be populated as restore runs (no up-front destination indexing). "
-        "First run may still do per-message duplicate checks; subsequent runs will skip quickly."
-    )
-
-    existing_dest_msg_ids_by_folder: Optional[dict[str, set[str]]] = {}  # lazily loaded per folder
-    existing_dest_msg_ids_lock: Optional[threading.Lock] = threading.Lock()
-
-    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
-        futures = []
-        for batch in batches:
-            # For Gmail, we use a special folder marker to indicate gmail-mode
-            # The process_restore_batch will determine the target folder per-email
-            # based on the manifest labels
-            futures.append(
-                executor.submit(
-                    process_restore_batch,
-                    batch,
-                    "__GMAIL_MODE__",  # Special marker - target determined per-email from manifest
-                    dest_conf,
-                    manifest,
-                    True,  # apply_labels
-                    apply_flags,  # apply_flags
-                    full_restore,
-                    existing_dest_msg_ids_by_folder,
-                    existing_dest_msg_ids_lock,
-                    progress_cache_data,
-                    progress_cache_lock,
-                    cache_path,
-                    dest_conf.get("host"),
-                    dest_conf.get("user"),
-                )
-            )
-
-        completed = 0
-        for future in concurrent.futures.as_completed(futures):
-            try:
-                future.result()
-                completed += 1
-                elapsed = time.time() - start_time
-                progress = (completed / len(batches)) * 100
-                safe_print(f"Progress: {progress:.1f}% ({elapsed:.0f}s elapsed)")
-            except Exception as e:
-                safe_print(f"Batch error: {e}")
-
-    elapsed = time.time() - start_time
-    safe_print(f"\nRestore completed in {elapsed:.1f}s")
-
-    # Force-flush progress cache at end.
-    restore_cache.maybe_save_dest_index_cache(cache_path, progress_cache_data, progress_cache_lock, force=True)
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Restore IMAP emails from local .eml files.")
-    parser.add_argument("--version", action="version", version=f"%(prog)s {imap_common.get_version()}")
-
-    # Source (Local Path)
-    env_path = os.getenv("BACKUP_LOCAL_PATH")
-    parser.add_argument(
-        "--src-path",
-        default=env_path,
-        required=not bool(env_path),
-        help="Local source path containing backup (or BACKUP_LOCAL_PATH)",
-    )
-
-    # Destination
-    default_dest_host = os.getenv("DEST_IMAP_HOST")
-    default_dest_user = os.getenv("DEST_IMAP_USERNAME")
-    default_dest_pass = os.getenv("DEST_IMAP_PASSWORD")
-    default_dest_client_id = os.getenv("DEST_OAUTH2_CLIENT_ID")
-
-    parser.add_argument(
-        "--dest-host",
-        default=default_dest_host,
-        required=not bool(default_dest_host),
-        help="Destination IMAP Server (or DEST_IMAP_HOST)",
-    )
-    parser.add_argument(
-        "--dest-user",
-        default=default_dest_user,
-        required=not bool(default_dest_user),
-        help="Destination Username (or DEST_IMAP_USERNAME)",
-    )
-
-    dest_auth_required = not bool(default_dest_pass or default_dest_client_id)
-    dest_auth = parser.add_mutually_exclusive_group(required=dest_auth_required)
-    dest_auth.add_argument(
-        "--dest-pass",
-        default=default_dest_pass,
-        help="Destination Password (or DEST_IMAP_PASSWORD)",
-    )
-    dest_auth.add_argument(
-        "--dest-oauth2-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help="Destination OAuth2 Client ID (or DEST_OAUTH2_CLIENT_ID)",
-    )
-    dest_auth.add_argument(
-        "--dest-client-id",
-        default=default_dest_client_id,
-        dest="dest_client_id",
-        help=argparse.SUPPRESS,
-    )
-    parser.add_argument(
-        "--dest-oauth2-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help="Destination OAuth2 Client Secret (if required) (or DEST_OAUTH2_CLIENT_SECRET)",
-    )
-    parser.add_argument(
-        "--dest-client-secret",
-        default=os.getenv("DEST_OAUTH2_CLIENT_SECRET"),
-        dest="dest_client_secret",
-        help=argparse.SUPPRESS,
-    )
-
-    # Config
-    parser.add_argument(
-        "--workers",
-        type=int,
-        default=int(os.getenv("MAX_WORKERS", 4)),
-        help="Thread count (default: 4)",
-    )
-    parser.add_argument(
-        "--batch",
-        type=int,
-        default=int(os.getenv("BATCH_SIZE", 10)),
-        help="Emails per batch",
-    )
-
-    # Gmail Labels
-    env_apply_labels = os.getenv("APPLY_LABELS", "false").lower() == "true"
-    parser.add_argument(
-        "--apply-labels",
-        action="store_true",
-        default=env_apply_labels,
-        help="Apply Gmail labels from labels_manifest.json",
-    )
-    env_apply_flags = os.getenv("APPLY_FLAGS", "false").lower() == "true"
-    parser.add_argument(
-        "--apply-flags",
-        action="store_true",
-        default=env_apply_flags,
-        help="Apply IMAP flags (read/starred/answered/draft) from manifest",
-    )
-    env_gmail_mode = os.getenv("GMAIL_MODE", "false").lower() == "true"
-    parser.add_argument(
-        "--gmail-mode",
-        action="store_true",
-        default=env_gmail_mode,
-        help="Gmail restore mode: Upload to INBOX and apply labels + flags from manifest",
-    )
-
-    env_full_restore = os.getenv("FULL_RESTORE", "false").lower() == "true"
-    parser.add_argument(
-        "--full-restore",
-        action="store_true",
-        default=env_full_restore,
-        help="Force full restore (legacy): process all emails and sync labels/flags for already-present messages.",
-    )
-
-    # Sync mode: delete from dest emails not in local backup
-    env_dest_delete = os.getenv("DEST_DELETE", "false").lower() == "true"
-    parser.add_argument(
-        "--dest-delete",
-        action="store_true",
-        default=env_dest_delete,
-        help="Delete emails from destination that don't exist in local backup (sync mode)",
-    )
-
-    # Optional folder filter
-    parser.add_argument("folder", nargs="?", help="Specific folder to restore")
-
-    args = parser.parse_args()
-
-    global MAX_WORKERS, BATCH_SIZE
-    MAX_WORKERS = args.workers
-    BATCH_SIZE = args.batch
-
-    # Build connection config (acquires OAuth2 token if configured)
-    dest_conf = imap_session.build_imap_conf(
-        args.dest_host, args.dest_user, args.dest_pass, args.dest_client_id, args.dest_client_secret, "destination"
-    )
-
-    # Expand path
-    local_path = os.path.expanduser(args.src_path)
-    if not os.path.exists(local_path):
-        print(f"Error: Source path does not exist: {local_path}")
-        sys.exit(1)
-
-    # Load manifest(s) if needed
-    manifest = {}
-    apply_labels = args.apply_labels or args.gmail_mode
-    apply_flags = args.apply_flags or args.gmail_mode
-
-    if apply_labels or apply_flags:
-        # Try to load labels manifest first (contains both labels and flags for Gmail backups)
-        manifest = imap_common.load_manifest(local_path, "labels_manifest.json")
-
-        # If no labels manifest, try flags-only manifest
-        if not manifest and apply_flags:
-            manifest = imap_common.load_manifest(local_path, "flags_manifest.json")
-
-        if not manifest:
-            print("Warning: No manifest found. Labels/flags will not be applied.")
-
-    progress_cache_file = None
-    progress_cache_data = None
-    progress_cache_lock = None
-    try:
-        progress_cache_file, progress_cache_data, progress_cache_lock = imap_common.load_progress_cache(
-            local_path,
-            args.dest_host,
-            args.dest_user,
-            log_fn=safe_print,
-        )
-    except Exception as e:
-        safe_print(f"Warning: Failed to load progress cache: {e}")
-
-    print("\n--- Configuration Summary ---")
-    print(f"Source Path     : {local_path}")
-    print(f"Destination Host: {args.dest_host}")
-    print(f"Destination User: {args.dest_user}")
-    print(f"Destination Auth: {imap_oauth2.auth_description(dest_conf['oauth2'] and dest_conf['oauth2']['provider'])}")
-    print(f"Workers         : {args.workers}")
-    if args.gmail_mode:
-        print("Mode            : Gmail Restore with Labels + Flags")
-    elif args.folder:
-        print(f"Target Folder   : {args.folder}")
-    if apply_labels:
-        print(f"Apply Labels    : Yes ({len(manifest)} mappings)")
-    if apply_flags:
-        print("Apply Flags     : Yes (read/starred/answered/draft)")
-    if args.dest_delete:
-        print("Dest Delete     : Yes (remove orphans from destination)")
-    print(
-        f"Restore Mode    : {'Full (all emails)' if args.full_restore else 'Incremental (new emails only, use --full-restore to restore all)'}"
-    )
-    print("-----------------------------\n")
-
-    try:
-        # Test connection
-        dest = imap_common.get_imap_connection_from_conf(dest_conf)
-        if not dest:
-            print("Error: Could not connect to destination server.")
-            sys.exit(1)
-
-        if args.gmail_mode:
-            dest.logout()
-            # Special Gmail mode
-            restore_gmail_with_labels(
-                local_path,
-                dest_conf,
-                manifest,
-                apply_flags,
-                full_restore=args.full_restore,
-                progress_cache_file=progress_cache_file,
-                progress_cache_data=progress_cache_data,
-                progress_cache_lock=progress_cache_lock,
-            )
-            dest = None  # Connection handled by restore_gmail_with_labels
-        elif args.folder:
-            # Restore specific folder
-            folder_path = os.path.join(local_path, args.folder.replace("/", os.sep))
-            if not os.path.exists(folder_path):
-                print(f"Error: Folder not found: {folder_path}")
-                sys.exit(1)
-            restore_folder(
-                args.folder,
-                folder_path,
-                dest_conf,
-                manifest,
-                apply_labels,
-                apply_flags,
-                args.dest_delete,
-                full_restore=args.full_restore,
-                cache_root=local_path,
-                progress_cache_file=progress_cache_file,
-                progress_cache_data=progress_cache_data,
-                progress_cache_lock=progress_cache_lock,
-            )
-            dest.logout()
-        else:
-            # Restore all folders
-            folders = imap_common.get_backup_folders(local_path)
-            if not folders:
-                print("No backup folders found.")
-                sys.exit(1)
-
-            print(f"Found {len(folders)} folders to restore.\n")
-            for folder_name, folder_path in folders:
-                # Skip manifest files
-                if folder_name in ("labels_manifest.json", "flags_manifest.json"):
-                    continue
-
-                # Proactively refresh OAuth2 token and ensure connection is healthy between folders
-                dest = imap_session.ensure_connection(dest, dest_conf)
-                if not dest:
-                    print("Fatal: Could not reconnect to destination IMAP server. Aborting.")
-                    sys.exit(1)
-
-                restore_folder(
-                    folder_name,
-                    folder_path,
-                    dest_conf,
-                    manifest,
-                    apply_labels,
-                    apply_flags,
-                    args.dest_delete,
-                    full_restore=args.full_restore,
-                    cache_root=local_path,
-                    progress_cache_file=progress_cache_file,
-                    progress_cache_data=progress_cache_data,
-                    progress_cache_lock=progress_cache_lock,
-                )
-
-            dest.logout()
-
-        print("\nRestore completed successfully.")
-
-    except KeyboardInterrupt:
-        raise
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    except KeyboardInterrupt:
-        print("\nProcess terminated by user.")
-        sys.exit(0)
-    except Exception as e:
-        print(f"Fatal Error: {e}")
-        import traceback
-
-        traceback.print_exc()
-        sys.exit(1)

src/imap_session.py

@@ -5,8 +5,8 @@ Connection and session management for IMAP operations with OAuth2 support.
 Combines imap_common (low-level IMAP) with imap_oauth2 (token refresh).
 """
 
-from auth import imap_oauth2
-from utils import imap_common
+import imap_common
+import imap_oauth2
 
 
 def build_imap_conf(host, user, password, client_id=None, client_secret=None, label=None):

src/provider_exchange.py

@@ -7,10 +7,9 @@ Constants and functions specific to Microsoft Exchange/Outlook IMAP implementati
 # Exchange/Outlook folders that typically can't be backed up via IMAP
 # These often contain proprietary data that Exchange returns as error messages
 EXCHANGE_SKIP_FOLDERS = {
-    "Calendar",
-    "Contacts",
-    "Conversation History",
     "Suggested Contacts",
+    "Conversation History",
+    "Calendar",
 }
 
 

src/provider_gmail.py

@@ -4,8 +4,8 @@ Gmail-Specific IMAP Utilities
 Constants and functions specific to Gmail/Google Workspace IMAP implementation.
 """
 
-from auth import imap_oauth2
-from utils import imap_common
+import imap_common
+import imap_oauth2
 
 # Gmail system folders
 GMAIL_ALL_MAIL = "[Gmail]/All Mail"

src/utils/imap_compress.py

@@ -1,180 +0,0 @@
-"""IMAP COMPRESS=DEFLATE support (RFC 4978).
-
-Provides a transparent compression wrapper for stdlib imaplib connections.
-After authentication, call ``enable_compression(conn)`` to negotiate
-COMPRESS=DEFLATE with the server.  All subsequent IMAP traffic is then
-compressed/decompressed automatically via zlib.
-
-The implementation follows imaplib's own STARTTLS pattern: send the
-command, and on OK replace ``self.sock`` with a wrapper.
-"""
-
-from __future__ import annotations
-
-import io
-import zlib
-from collections.abc import Callable
-
-
-class _CompressedRawIO(io.RawIOBase):
-    """Minimal RawIOBase adapter so ``makefile()`` can return a BufferedReader
-    that reads through the decompression layer."""
-
-    def __init__(self, compressed_sock: _CompressedSocket):
-        self._sock = compressed_sock
-
-    def readable(self) -> bool:
-        return True
-
-    def readinto(self, b: bytearray | memoryview) -> int:
-        data = self._sock.recv(len(b))
-        if not data:
-            return 0
-        n = len(data)
-        b[:n] = data
-        return n
-
-
-class _CompressedSocket:
-    """Socket wrapper that adds zlib DEFLATE compression/decompression.
-
-    Wraps a real socket (typically SSL) so that:
-    - ``sendall()`` compresses data before sending
-    - ``recv()`` decompresses data after receiving
-
-    Python 3.13's imaplib reads via ``self.sock.recv()`` directly,
-    so intercepting at the socket level is sufficient.
-    """
-
-    def __init__(self, sock, initial_compressed: bytes = b""):
-        self._sock = sock
-        self._compressor = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        self._decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
-        # Buffer for decompressed data not yet consumed by recv()
-        self._recv_buf = b""
-        # If there was read-ahead data in imaplib's buffer after the
-        # COMPRESS OK response, it is already compressed and must be
-        # fed to the decompressor.
-        if initial_compressed:
-            self._recv_buf = self._decompressor.decompress(initial_compressed)
-
-    # -- outgoing (compress) ---------------------------------------------------
-
-    def sendall(self, data: bytes) -> None:
-        compressed = self._compressor.compress(data)
-        compressed += self._compressor.flush(zlib.Z_SYNC_FLUSH)
-        self._sock.sendall(compressed)
-
-    def send(self, data: bytes) -> int:
-        self.sendall(data)
-        return len(data)
-
-    # -- incoming (decompress) -------------------------------------------------
-
-    def recv(self, bufsize: int) -> bytes:
-        # Return buffered decompressed data first
-        if self._recv_buf:
-            chunk = self._recv_buf[:bufsize]
-            self._recv_buf = self._recv_buf[bufsize:]
-            return chunk
-
-        # Read and decompress until we have at least 1 byte or true EOF.
-        # zlib can produce b"" for partial DEFLATE blocks; returning that
-        # would signal EOF to imaplib and break the connection.
-        decompressed = b""
-        while not decompressed:
-            raw = self._sock.recv(max(bufsize, 16384))
-            if not raw:
-                return b""  # real EOF
-            decompressed = self._decompressor.decompress(raw)
-
-        if len(decompressed) > bufsize:
-            self._recv_buf = decompressed[bufsize:]
-            return decompressed[:bufsize]
-        return decompressed
-
-    # -- lifecycle / proxy -----------------------------------------------------
-
-    def shutdown(self, how: int) -> None:
-        self._sock.shutdown(how)
-
-    def close(self) -> None:
-        self._sock.close()
-
-    def makefile(self, *args, **kwargs):
-        # Older Python imaplib reads from conn._file (a file object)
-        # rather than conn.sock.recv() directly.  The file must read
-        # through the decompression layer, not from the raw socket.
-        return io.BufferedReader(_CompressedRawIO(self))
-
-    def __getattr__(self, name: str):
-        return getattr(self._sock, name)
-
-
-def enable_compression(
-    conn,
-    *,
-    log_fn: Callable[[str], None] | None = None,
-) -> bool:
-    """Negotiate COMPRESS=DEFLATE on a raw imaplib IMAP4/IMAP4_SSL connection.
-
-    Must be called **after** authentication (login/XOAUTH2).
-
-    Returns True if compression was successfully enabled, False otherwise
-    (e.g. server does not advertise the capability, or the command failed).
-    """
-    raw_caps = getattr(conn, "capabilities", ())
-    caps = {c.decode() if isinstance(c, bytes) else c for c in raw_caps}
-    if "COMPRESS=DEFLATE" not in caps:
-        if log_fn is not None:
-            log_fn("Deflate compression not supported by server")
-        return False
-
-    try:
-        typ, _data = conn._simple_command("COMPRESS", "DEFLATE")
-    except Exception:
-        if log_fn is not None:
-            log_fn("Deflate compression negotiation failed")
-        return False
-
-    if typ != "OK":
-        if log_fn is not None:
-            log_fn("Deflate compression rejected by server")
-        return False
-
-    # Drain any read-ahead data from imaplib's internal buffers.
-    # After the server sends OK, all subsequent data is compressed.
-    # If recv() read past the OK response, the leftover bytes are
-    # already compressed and must be fed to the decompressor.
-    leftover = b""
-
-    # conn._file (BufferedReader) may have read-ahead bytes beyond
-    # the OK response that are already compressed.  Drain them.
-    old_file = getattr(conn, "_file", None)
-    if old_file is not None:
-        try:
-            buffered = old_file.peek(65536)
-            if isinstance(buffered, (bytes, bytearray)) and buffered:
-                leftover += bytes(buffered)
-        except Exception:
-            pass
-
-    readbuf = getattr(conn, "_readbuf", None)
-    if readbuf:
-        if isinstance(readbuf, (bytes, bytearray)):
-            leftover += bytes(readbuf)
-        else:
-            leftover += b"".join(readbuf)
-        if isinstance(readbuf, bytes):
-            conn._readbuf = b""
-        else:
-            readbuf.clear()
-
-    # Replace the socket (follows the STARTTLS pattern in imaplib).
-    conn.sock = _CompressedSocket(conn.sock, initial_compressed=leftover)
-    conn._file = conn.sock.makefile("rb")
-
-    if log_fn is not None:
-        log_fn("Deflate compression enabled")
-
-    return True

test/test_backup_imap_emails.py

@@ -17,10 +17,10 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_backup as backup_imap_emails
+import backup_imap_emails
+import imap_common
+import provider_gmail
 from conftest import temp_argv, temp_env
-from providers import provider_gmail
-from utils import imap_common
 
 
 def _mock_imap_env(port):

test/test_compare_imap_folders.py

@@ -18,7 +18,7 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_compare as compare_imap_folders
+import compare_imap_folders
 from conftest import temp_argv, temp_env
 
 

test/test_count_imap_emails.py

@@ -16,9 +16,9 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_count as count_imap_emails
+import count_imap_emails
+import imap_common
 from conftest import temp_argv, temp_env
-from utils import imap_common
 
 
 def _mock_imap_env(port):

test/test_imap_common.py

@@ -18,10 +18,10 @@ from unittest.mock import Mock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
+import imap_common
 from conftest import temp_env
-from utils import imap_common
 
 
 class TestVerifyEnvVars:
@@ -320,6 +320,48 @@ class TestDetectTrashFolder:
         assert result is None
 
 
+class TestMessageExistsInFolder:
+    """Tests for message_exists_in_folder function."""
+
+    def test_no_message_id(self):
+        """Test returns False when message_id is None."""
+        mock_conn = Mock()
+        result = imap_common.message_exists_in_folder(mock_conn, None)
+        assert result is False
+
+    def test_search_fails(self):
+        """Test returns False when search fails."""
+        mock_conn = Mock()
+        mock_conn.search.return_value = ("NO", [])
+
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
+        assert result is False
+
+    def test_no_matches(self):
+        """Test returns False when no matches found."""
+        mock_conn = Mock()
+        mock_conn.search.return_value = ("OK", [b""])
+
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
+        assert result is False
+
+    def test_match_found(self):
+        """Test returns True when message with same ID found."""
+        mock_conn = Mock()
+        mock_conn.search.return_value = ("OK", [b"1"])
+
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
+        assert result is True
+
+    def test_search_exception(self):
+        """Test returns False when search raises an exception."""
+        mock_conn = Mock()
+        mock_conn.search.side_effect = Exception("Connection error")
+
+        result = imap_common.message_exists_in_folder(mock_conn, "<msg-id>")
+        assert result is False
+
+
 class TestExtractMessageId:
     """Tests for extract_message_id function."""
 
@@ -765,77 +807,3 @@ class TestGetMessageIdsInFolder:
 
         result = imap_common.get_message_ids_in_folder(mock_conn)
         assert set(result.values()) == {"<valid@example.com>"}
-
-
-class TestLoadFolderMsgIds:
-    """Tests for load_folder_msg_ids() using the mock IMAP server."""
-
-    def test_returns_none_when_disabled(self):
-        """Returns None if dict or lock is None (tracking disabled)."""
-        mock_conn = Mock()
-        import threading
-
-        lock = threading.Lock()
-        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, lock) is None
-        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", {}, None) is None
-        assert imap_common.load_folder_msg_ids(mock_conn, "INBOX", None, None) is None
-
-    def test_fetches_message_ids_from_server(self, single_mock_server):
-        """Fetches Message-IDs from server and caches them."""
-        import imaplib
-        import threading
-
-        server_data = {
-            "TestFolder": [
-                b"Subject: A\r\nMessage-ID: <a@test.com>\r\n\r\nBody A",
-                b"Subject: B\r\nMessage-ID: <b@test.com>\r\n\r\nBody B",
-            ]
-        }
-        _server, port = single_mock_server(server_data)
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        lock = threading.Lock()
-        by_folder: dict[str, set[str]] = {}
-
-        result = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
-        assert result == {"<a@test.com>", "<b@test.com>"}
-        assert "TestFolder" in by_folder
-
-        # Second call returns the same cached set object
-        result2 = imap_common.load_folder_msg_ids(conn, "TestFolder", by_folder, lock)
-        assert result2 is result
-        conn.logout()
-
-    def test_empty_folder(self, single_mock_server):
-        """Returns empty set for a folder with no messages."""
-        import imaplib
-        import threading
-
-        _server, port = single_mock_server({"INBOX": []})
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        lock = threading.Lock()
-        by_folder: dict[str, set[str]] = {}
-
-        result = imap_common.load_folder_msg_ids(conn, "INBOX", by_folder, lock)
-        assert result == set()
-        conn.logout()
-
-    def test_creates_folder_if_missing(self, single_mock_server):
-        """Creates folder on first access if it doesn't exist."""
-        import imaplib
-        import threading
-
-        server, port = single_mock_server({"INBOX": []})
-        conn = imaplib.IMAP4("localhost", port)
-        conn.login("user", "pass")
-
-        lock = threading.Lock()
-        by_folder: dict[str, set[str]] = {}
-
-        result = imap_common.load_folder_msg_ids(conn, "NewFolder", by_folder, lock)
-        assert result == set()
-        assert "NewFolder" in server.folders
-        conn.logout()

test/test_imap_oauth2.py

@@ -14,9 +14,11 @@ from unittest.mock import patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from auth import imap_oauth2, oauth2_google, oauth2_microsoft
+import imap_oauth2
+import oauth2_google
+import oauth2_microsoft
 
 
 @pytest.fixture(autouse=True)

test/test_imap_retry.py

@@ -18,8 +18,8 @@ import pytest
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../tools")))
 
+import imap_retry
 from mock_imap_server import start_server_thread as start_mock_server
-from utils import imap_retry
 
 
 class TestConnectionProxy:

test/test_imap_session.py

@@ -12,9 +12,9 @@ import os
 import sys
 from unittest.mock import MagicMock, patch
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from core import imap_session
+import imap_session
 
 
 class TestBuildImapConf:

test/test_migrate_imap_emails.py

@@ -19,9 +19,9 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_migrate as migrate_imap_emails
+import imap_common
+import migrate_imap_emails
 from conftest import temp_argv, temp_env
-from utils import imap_common
 
 
 def _mock_migrate_env(src_port, dest_port):
@@ -285,7 +285,7 @@ class TestCacheHitWithoutLock:
 
         src.select('"INBOX"', readonly=False)
 
-        existing_dest_msg_ids_by_folder = {"INBOX": {msg_id}}
+        existing_dest_msg_ids = {msg_id}
         success, _src, _dest, deleted = migrate_imap_emails.process_single_uid(
             src,
             dest,
@@ -298,7 +298,7 @@ class TestCacheHitWithoutLock:
             None,
             True,
             False,
-            existing_dest_msg_ids_by_folder=existing_dest_msg_ids_by_folder,
+            existing_dest_msg_ids=existing_dest_msg_ids,
             existing_dest_msg_ids_lock=None,
             progress_cache_path=None,
             progress_cache_data=None,
@@ -357,18 +357,16 @@ class TestMigrateErrorHandling:
 
         from unittest.mock import patch
 
-        original_get_imap_connection = imap_common.get_imap_connection
-
         def side_effect(host, user, pwd, oauth2_token=None):
             if threading.current_thread() is threading.main_thread():
-                return original_get_imap_connection(host, user, pwd, oauth2_token)
+                return imap_common.get_imap_connection(host, user, pwd, oauth2_token)
             return None  # Simulate connection failure in worker
 
         env = _mock_migrate_env(p1, p2)
         with (
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
-            patch("utils.imap_common.get_imap_connection", side_effect=side_effect),
+            patch("imap_common.get_imap_connection", side_effect=side_effect),
         ):
             # Should not crash, but log error
             migrate_imap_emails.main()
@@ -444,7 +442,7 @@ class TestMigrateErrorHandling:
 
         env = _mock_migrate_env(p1, p2)
         with (
-            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid, autospec=True),
+            patch.object(imaplib.IMAP4, "uid", side_effect=side_effect_uid),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
         ):
@@ -465,7 +463,7 @@ class TestMigrateErrorHandling:
             "DEST_IMAP_USERNAME": "dest_user",
             "DEST_IMAP_PASSWORD": "p",
         }
-        with patch("utils.imap_common.get_imap_connection", return_value=None), temp_env(env):
+        with patch("imap_common.get_imap_connection", return_value=None), temp_env(env):
             with pytest.raises(SystemExit) as exc:
                 migrate_imap_emails.main()
         assert exc.value.code == 1
@@ -482,7 +480,7 @@ class TestMigrateErrorHandling:
         env = _mock_migrate_env(p1, p2)
         with (
             patch(
-                "utils.imap_common.load_progress_cache",
+                "imap_common.load_progress_cache",
                 lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
             ),
             temp_env(env),
@@ -509,7 +507,7 @@ class TestTrashHandling:
         env = _mock_migrate_env(p1, p2)
         env["DELETE_FROM_SOURCE"] = "true"
         with (
-            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
+            patch("imap_common.detect_trash_folder", return_value="Trash"),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py"]),
         ):
@@ -531,7 +529,7 @@ class TestTrashHandling:
         env = _mock_migrate_env(p1, p2)
         env["DELETE_FROM_SOURCE"] = "true"
         with (
-            patch("utils.imap_common.detect_trash_folder", return_value="Trash"),
+            patch("imap_common.detect_trash_folder", return_value="Trash"),
             temp_env(env),
             temp_argv(["migrate_imap_emails.py", "INBOX"]),
         ):

test/test_migrate_resume.py

@@ -6,9 +6,9 @@ import sys
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_migrate as migrate_imap_emails
+import migrate_imap_emails
+import restore_cache
 from conftest import temp_argv, temp_env
-from utils import restore_cache
 
 
 def _run_migrate(cache_dir, src_port, dest_port):

test/test_migrate_with_cache.py

@@ -10,10 +10,11 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_migrate as migrate_imap_emails
+import imap_common
+import imap_session
+import migrate_imap_emails
+import restore_cache
 from conftest import make_mock_connection, temp_argv, temp_env
-from core import imap_session
-from utils import imap_common, restore_cache
 
 
 def _run_migrate(cache_dir, src_port, dest_port, full_migrate=False, extra_env=None):
@@ -167,7 +168,7 @@ class TestMigrationCache:
         dest.login("dest", "p")
 
         with (
-            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch("imap_common.get_imap_connection", make_mock_connection(p1, p2)),
             patch.object(
                 imap_common,
                 "load_progress_cache",
@@ -214,7 +215,7 @@ class TestMigrationCache:
         dest.login("dest", "p")
 
         with (
-            patch("utils.imap_common.get_imap_connection", make_mock_connection(p1, p2)),
+            patch("imap_common.get_imap_connection", make_mock_connection(p1, p2)),
             patch.object(
                 restore_cache,
                 "get_cached_message_ids",

test/test_oauth2_google.py

@@ -12,9 +12,9 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from auth import oauth2_google
+import oauth2_google
 
 
 @pytest.fixture(autouse=True)

test/test_oauth2_microsoft.py

@@ -14,9 +14,9 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from auth import oauth2_microsoft
+import oauth2_microsoft
 from conftest import temp_env
 from mock_oauth_server import MOCK_TENANT_ID
 

test/test_provider_exchange.py

@@ -9,9 +9,9 @@ Tests cover:
 import os
 import sys
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from providers import provider_exchange
+import provider_exchange
 
 
 class TestIsSpecialFolder:
@@ -29,10 +29,6 @@ class TestIsSpecialFolder:
         """Test that Calendar is detected as special folder."""
         assert provider_exchange.is_special_folder("Calendar") is True
 
-    def test_contacts_is_special(self):
-        """Test that Contacts is detected as special folder."""
-        assert provider_exchange.is_special_folder("Contacts") is True
-
     def test_inbox_is_not_special(self):
         """Test that INBOX is not a special folder."""
         assert provider_exchange.is_special_folder("INBOX") is False
@@ -80,12 +76,11 @@ class TestExchangeConstants:
         assert "Suggested Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
         assert "Conversation History" in provider_exchange.EXCHANGE_SKIP_FOLDERS
         assert "Calendar" in provider_exchange.EXCHANGE_SKIP_FOLDERS
-        assert "Contacts" in provider_exchange.EXCHANGE_SKIP_FOLDERS
 
     def test_exchange_skip_folders_count(self):
         """Test that EXCHANGE_SKIP_FOLDERS has expected number of entries."""
         # This test will catch if folders are accidentally added/removed
-        assert len(provider_exchange.EXCHANGE_SKIP_FOLDERS) == 4
+        assert len(provider_exchange.EXCHANGE_SKIP_FOLDERS) == 3
 
     def test_inbox_not_in_skip_folders(self):
         """Test that INBOX is not in EXCHANGE_SKIP_FOLDERS."""

test/test_provider_gmail.py

@@ -12,9 +12,9 @@ import os
 import sys
 from unittest.mock import MagicMock, patch
 
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src")))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-from providers import provider_gmail
+import provider_gmail
 
 
 class TestIsLabelFolder:
@@ -181,7 +181,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        from utils import imap_common
+        import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),
@@ -221,7 +221,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        from utils import imap_common
+        import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),
@@ -251,7 +251,7 @@ class TestBuildGmailLabelIndex:
         def mock_safe_print(msg):
             pass
 
-        from utils import imap_common
+        import imap_common
 
         with (
             patch.object(imap_common, "list_selectable_folders", mock_list_folders),

test/test_restore_imap_emails.py

@@ -18,9 +18,10 @@ import pytest
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../src")))
 
-import imap_restore as restore_imap_emails
+import imap_common
+import restore_cache
+import restore_imap_emails
 from conftest import temp_argv, temp_env
-from utils import imap_common, restore_cache
 
 
 def _mock_restore_env(port):
@@ -508,12 +509,45 @@ class TestRestoreE2EHelpers:
             "INBOX",
             b"Subject: Upload\r\nMessage-ID: <up@test>\r\n\r\nBody",
             '"15-Jan-2024 10:30:00 +0000"',
+            "<up@test>",
         )
 
         assert result == restore_imap_emails.UploadResult.SUCCESS
         assert len(server.folders["INBOX"]) == 1
         conn.logout()
 
+    def test_upload_email_duplicate(self, single_mock_server):
+        dest_data = {"INBOX": [b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody"]}
+        server, port = single_mock_server(dest_data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+
+        result = restore_imap_emails.upload_email(
+            conn,
+            "INBOX",
+            b"Subject: Dup\r\nMessage-ID: <dup@test>\r\n\r\nBody",
+            '"15-Jan-2024 10:30:00 +0000"',
+            "<dup@test>",
+            check_duplicate=True,
+        )
+
+        assert result == restore_imap_emails.UploadResult.ALREADY_EXISTS
+        assert len(server.folders["INBOX"]) == 1
+        conn.logout()
+
+    def test_email_exists_in_folder(self, single_mock_server):
+        dest_data = {"INBOX": [b"Subject: Exists\r\nMessage-ID: <exists@test>\r\n\r\nBody"]}
+        _server, port = single_mock_server(dest_data)
+
+        conn = imaplib.IMAP4("localhost", port)
+        conn.login("user", "pass")
+        conn.select('"INBOX"')
+
+        assert restore_imap_emails.email_exists_in_folder(conn, "<exists@test>") is True
+        assert restore_imap_emails.email_exists_in_folder(conn, "<missing@test>") is False
+        conn.logout()
+
     def test_sync_flags_on_existing(self, single_mock_server):
         dest_data = {
             "INBOX": [{"uid": 1, "flags": set(), "content": b"Subject: Flag\r\nMessage-ID: <flag@test>\r\n\r\nBody"}]
@@ -753,8 +787,8 @@ Body content.
 
         # Track calls to record_progress
         with (
-            patch("utils.restore_cache.record_progress") as mock_record_progress,
-            patch("utils.imap_common.get_imap_connection", return_value=mock_conn),
+            patch("restore_cache.record_progress") as mock_record_progress,
+            patch("imap_common.get_imap_connection", return_value=mock_conn),
         ):
             env = {
                 "DEST_IMAP_HOST": "localhost",

test/utils/test_imap_compress.py

@@ -1,394 +0,0 @@
-"""Tests for IMAP COMPRESS=DEFLATE support."""
-
-import zlib
-from unittest.mock import MagicMock, patch
-
-from utils import imap_compress
-
-
-class TestCompressedSocket:
-    """Tests for the _CompressedSocket wrapper."""
-
-    def _make_compressed(self, plaintext: bytes) -> bytes:
-        """Compress data using raw DEFLATE with Z_SYNC_FLUSH (server-side)."""
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        return c.compress(plaintext) + c.flush(zlib.Z_SYNC_FLUSH)
-
-    def test_sendall_compresses_data(self):
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock)
-
-        ws.sendall(b"AAAA0 SELECT INBOX\r\n")
-
-        sock.sendall.assert_called_once()
-        compressed = sock.sendall.call_args[0][0]
-        # Verify round-trip: decompress the sent data
-        d = zlib.decompressobj(-zlib.MAX_WBITS)
-        assert d.decompress(compressed) == b"AAAA0 SELECT INBOX\r\n"
-
-    def test_send_compresses_and_returns_length(self):
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock)
-
-        result = ws.send(b"hello")
-        assert result == 5
-        sock.sendall.assert_called_once()
-
-    def test_recv_decompresses_data(self):
-        plaintext = b"* OK IMAP ready\r\n"
-        compressed = self._make_compressed(plaintext)
-
-        sock = MagicMock()
-        sock.recv.return_value = compressed
-
-        ws = imap_compress._CompressedSocket(sock)
-        result = ws.recv(4096)
-
-        assert result == plaintext
-
-    def test_recv_buffers_excess(self):
-        plaintext = b"A" * 100
-        compressed = self._make_compressed(plaintext)
-
-        sock = MagicMock()
-        sock.recv.return_value = compressed
-
-        ws = imap_compress._CompressedSocket(sock)
-        # Request only 30 bytes — rest should be buffered
-        chunk1 = ws.recv(30)
-        assert chunk1 == b"A" * 30
-        assert len(ws._recv_buf) == 70
-
-        # Next recv should return from buffer without hitting socket
-        chunk2 = ws.recv(70)
-        assert chunk2 == b"A" * 70
-        sock.recv.assert_called_once()  # Only one real socket read
-
-    def test_recv_empty_returns_empty(self):
-        sock = MagicMock()
-        sock.recv.return_value = b""
-
-        ws = imap_compress._CompressedSocket(sock)
-        assert ws.recv(4096) == b""
-
-    def test_initial_compressed_data_decompressed(self):
-        plaintext = b"* 5 EXISTS\r\n"
-        compressed = self._make_compressed(plaintext)
-
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock, initial_compressed=compressed)
-
-        # Should return the pre-decompressed data without hitting socket
-        result = ws.recv(4096)
-        assert result == plaintext
-        sock.recv.assert_not_called()
-
-    def test_shutdown_proxied(self):
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock)
-        ws.shutdown(2)
-        sock.shutdown.assert_called_once_with(2)
-
-    def test_close_proxied(self):
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock)
-        ws.close()
-        sock.close.assert_called_once()
-
-    def test_makefile_returns_buffered_reader(self):
-        import io
-
-        plaintext = b"* OK ready\r\n"
-        compressed = self._make_compressed(plaintext)
-
-        sock = MagicMock()
-        sock.recv.return_value = compressed
-
-        ws = imap_compress._CompressedSocket(sock)
-        f = ws.makefile("rb")
-        assert isinstance(f, io.BufferedReader)
-        assert f.readline() == plaintext
-
-    def test_getattr_proxied(self):
-        sock = MagicMock()
-        sock.gettimeout.return_value = 30
-        ws = imap_compress._CompressedSocket(sock)
-        assert ws.gettimeout() == 30
-
-    def test_multiple_sendall_share_compressor_state(self):
-        """Verify the compressor maintains state across calls (streaming)."""
-        sock = MagicMock()
-        ws = imap_compress._CompressedSocket(sock)
-
-        ws.sendall(b"first ")
-        ws.sendall(b"second")
-
-        assert sock.sendall.call_count == 2
-        # Both calls should produce decompressible data
-        d = zlib.decompressobj(-zlib.MAX_WBITS)
-        result = b""
-        for call in sock.sendall.call_args_list:
-            result += d.decompress(call[0][0])
-        assert result == b"first second"
-
-    def test_recv_loops_on_partial_deflate_block(self):
-        """Decompress can return b'' for partial blocks; recv must not return that as EOF."""
-        sock = MagicMock()
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-
-        # Build a compressed stream, then split it mid-block so the first
-        # chunk decompresses to b"" and the second completes the block.
-        full_compressed = c.compress(b"hello\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
-        # Split at byte 1 — the first fragment is too short to produce output
-        part1 = full_compressed[:1]
-        part2 = full_compressed[1:]
-        sock.recv.side_effect = [part1, part2]
-
-        ws = imap_compress._CompressedSocket(sock)
-        result = ws.recv(4096)
-
-        assert result == b"hello\r\n"
-        assert sock.recv.call_count == 2
-
-    def test_recv_returns_empty_on_real_eof(self):
-        """True socket EOF (b'') should still return b''."""
-        sock = MagicMock()
-        sock.recv.return_value = b""
-
-        ws = imap_compress._CompressedSocket(sock)
-        assert ws.recv(4096) == b""
-
-    def test_recv_multiple_chunks(self):
-        """Simulate receiving data across multiple socket reads."""
-        chunk1 = b"line one\r\n"
-        chunk2 = b"line two\r\n"
-
-        sock = MagicMock()
-        # Server compressor maintains state across writes
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        compressed1 = c.compress(chunk1) + c.flush(zlib.Z_SYNC_FLUSH)
-        compressed2 = c.compress(chunk2) + c.flush(zlib.Z_SYNC_FLUSH)
-        sock.recv.side_effect = [compressed1, compressed2]
-
-        ws = imap_compress._CompressedSocket(sock)
-        assert ws.recv(4096) == chunk1
-        assert ws.recv(4096) == chunk2
-
-
-class TestEnableCompression:
-    """Tests for the enable_compression() function."""
-
-    def test_returns_false_when_capability_missing(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "IDLE")
-        assert imap_compress.enable_compression(conn) is False
-
-    def test_returns_false_when_capability_missing_bytes(self):
-        conn = MagicMock()
-        conn.capabilities = (b"IMAP4REV1", b"IDLE")
-        assert imap_compress.enable_compression(conn) is False
-
-    def test_returns_false_when_no_capabilities_attr(self):
-        conn = MagicMock(spec=[])
-        assert imap_compress.enable_compression(conn) is False
-
-    def test_returns_false_when_command_fails(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("NO", [b"compression not available"])
-
-        assert imap_compress.enable_compression(conn) is False
-
-    def test_returns_false_when_command_raises(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.side_effect = Exception("network error")
-
-        assert imap_compress.enable_compression(conn) is False
-
-    def test_enables_compression_with_bytes_capabilities(self):
-        """Real imaplib stores capabilities as bytes; verify they are handled."""
-        conn = MagicMock()
-        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
-        conn._readbuf = []
-        conn.sock = MagicMock()
-
-        result = imap_compress.enable_compression(conn)
-
-        assert result is True
-        assert isinstance(conn.sock, imap_compress._CompressedSocket)
-
-    def test_enables_compression_on_success(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"DEFLATE active"])
-        conn._readbuf = []
-        conn.sock = MagicMock()
-
-        result = imap_compress.enable_compression(conn)
-
-        assert result is True
-        assert isinstance(conn.sock, imap_compress._CompressedSocket)
-
-    def test_drains_readbuf_on_success(self):
-        # Simulate leftover compressed data in the read buffer
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        compressed = c.compress(b"* 1 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
-
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = [compressed]
-        conn.sock = MagicMock()
-
-        imap_compress.enable_compression(conn)
-
-        # readbuf should have been cleared
-        assert len(conn._readbuf) == 0
-        # The compressed leftover should be pre-decompressed in the wrapper
-        assert conn.sock._recv_buf == b"* 1 EXISTS\r\n"
-
-    def test_drains_bytearray_readbuf(self):
-        """Real imaplib uses bytearray for _readbuf; verify drain handles it."""
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        compressed = c.compress(b"* 2 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
-
-        conn = MagicMock()
-        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = bytearray(compressed)
-        conn.sock = MagicMock()
-
-        imap_compress.enable_compression(conn)
-
-        assert len(conn._readbuf) == 0
-        assert conn.sock._recv_buf == b"* 2 EXISTS\r\n"
-
-    def test_drains_bytes_readbuf(self):
-        """Some imaplib versions use bytes for _readbuf; verify drain handles it."""
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        compressed = c.compress(b"* 3 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
-
-        conn = MagicMock()
-        conn.capabilities = (b"IMAP4REV1", b"COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = bytes(compressed)
-        conn.sock = MagicMock()
-
-        imap_compress.enable_compression(conn)
-
-        assert conn._readbuf == b""
-        assert conn.sock._recv_buf == b"* 3 EXISTS\r\n"
-
-    def test_drains_file_buffer(self):
-        """BufferedReader may have read-ahead bytes beyond the OK response."""
-        import io
-
-        c = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, -zlib.MAX_WBITS)
-        compressed = c.compress(b"* 4 EXISTS\r\n") + c.flush(zlib.Z_SYNC_FLUSH)
-
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = b""
-        conn.sock = MagicMock()
-
-        # Simulate a BufferedReader with read-ahead compressed bytes
-        old_file = io.BufferedReader(io.BytesIO(compressed))
-        old_file.peek(65536)  # fill the buffer
-        conn._file = old_file
-
-        imap_compress.enable_compression(conn)
-
-        assert conn.sock._recv_buf == b"* 4 EXISTS\r\n"
-
-    def test_calls_log_fn(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = []
-        conn.sock = MagicMock()
-
-        messages = []
-        imap_compress.enable_compression(conn, log_fn=messages.append)
-
-        assert len(messages) == 1
-        assert "compression enabled" in messages[0]
-
-    def test_logs_when_not_supported(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1",)  # No COMPRESS
-
-        messages = []
-        imap_compress.enable_compression(conn, log_fn=messages.append)
-
-        assert len(messages) == 1
-        assert "not supported" in messages[0]
-
-    def test_logs_when_rejected(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("NO", [b"not available"])
-
-        messages = []
-        imap_compress.enable_compression(conn, log_fn=messages.append)
-
-        assert len(messages) == 1
-        assert "rejected" in messages[0]
-
-    def test_logs_when_negotiation_fails(self):
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.side_effect = Exception("network error")
-
-        messages = []
-        imap_compress.enable_compression(conn, log_fn=messages.append)
-
-        assert len(messages) == 1
-        assert "failed" in messages[0]
-
-    def test_recreates_file(self):
-        import io
-
-        conn = MagicMock()
-        conn.capabilities = ("IMAP4REV1", "COMPRESS=DEFLATE")
-        conn._simple_command.return_value = ("OK", [b"ok"])
-        conn._readbuf = []
-        real_sock = MagicMock()
-        conn.sock = real_sock
-
-        imap_compress.enable_compression(conn)
-
-        # _file should be a BufferedReader over the compressed socket
-        assert isinstance(conn._file, io.BufferedReader)
-
-
-class TestIntegrationWithGetImapConnection:
-    """Test that enable_compression is called during connection setup."""
-
-    @patch("utils.imap_common.imaplib")
-    @patch("utils.imap_compress.enable_compression")
-    def test_compression_attempted_after_login(self, mock_compress, mock_imaplib):
-        from utils import imap_common
-
-        mock_conn = MagicMock()
-        mock_imaplib.IMAP4_SSL.return_value = mock_conn
-
-        imap_common.get_imap_connection("imap.example.com", "user", password="pass")
-
-        mock_conn.login.assert_called_once()
-        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)
-
-    @patch("utils.imap_common.imaplib")
-    @patch("utils.imap_compress.enable_compression")
-    def test_compression_attempted_after_oauth2(self, mock_compress, mock_imaplib):
-        from utils import imap_common
-
-        mock_conn = MagicMock()
-        mock_imaplib.IMAP4_SSL.return_value = mock_conn
-
-        imap_common.get_imap_connection("imap.example.com", "user", oauth2_token="token123")
-
-        mock_conn.authenticate.assert_called_once()
-        mock_compress.assert_called_once_with(mock_conn, log_fn=imap_common.safe_print)