#!/usr/bin/env python3
"""
Proton Bridge → Stalwart IMAP migration
"""
import imaplib, ssl, email, time, sys, traceback
from datetime import datetime

SRC_HOST = '127.0.0.1'; SRC_PORT = 1143
SRC_USER = 'tj@jongsma.me'; SRC_PASS = 'BlcMCKtNDfqv0cq1LmGR9g'
DST_HOST = 'mail.jongsma.me'; DST_PORT = 993
DST_USER = 'tj'; DST_PASS = 'TJ-Migrate-2026!'

SKIP_FOLDERS = {'All Mail', 'Starred', 'Labels/Sent Messages',
                'Labels/Deleted Messages', 'Labels/Snoozed'}
SKIP_PREFIXES = ('Labels/',)

BATCH = 10  # smaller batch = more reliable with Proton Bridge

def src_ctx():
    c = ssl.create_default_context()
    c.check_hostname = False; c.verify_mode = ssl.CERT_NONE
    return c

def connect_src():
    s = imaplib.IMAP4(SRC_HOST, SRC_PORT)
    s.starttls(src_ctx()); s.login(SRC_USER, SRC_PASS)
    return s

def connect_dst():
    s = imaplib.IMAP4_SSL(DST_HOST, DST_PORT, ssl_context=ssl.create_default_context())
    s.login(DST_USER, DST_PASS)
    return s

def ensure_dst_folder(dst, folder):
    st, _ = dst.select(folder)
    if st == 'OK': return
    dst.create(folder)
    dst.select(folder)

def get_dst_message_ids(dst, folder):
    ids = set()
    st, _ = dst.select(folder)
    if st != 'OK': return ids
    st2, data = dst.search(None, 'ALL')
    if st2 != 'OK' or not data[0]: return ids
    msgs = data[0].split()
    for i in range(0, len(msgs), 200):
        chunk = b','.join(msgs[i:i+200])
        try:
            st3, hd = dst.fetch(chunk, '(BODY[HEADER.FIELDS (MESSAGE-ID)])')
            for item in hd:
                if isinstance(item, tuple):
                    raw = item[1].decode('utf-8', errors='replace')
                    for line in raw.splitlines():
                        if line.lower().startswith('message-id:'):
                            ids.add(line.split(':',1)[1].strip())
        except: pass
    return ids

def sync_folder(src, dst, folder):
    copied = skipped = errors = 0
    src_q = f'"{folder}"' if ' ' in folder or '/' in folder else folder

    st, _ = src.select(src_q, readonly=True)
    if st != 'OK':
        print(f'  ✗ Cannot select src: {folder}'); return 0, 0, 1

    st2, data = src.search(None, 'ALL')
    if st2 != 'OK' or not data[0]: return 0, 0, 0
    msg_ids = data[0].split()
    src_count = len(msg_ids)

    ensure_dst_folder(dst, folder)
    existing_ids = get_dst_message_ids(dst, folder)
    dst.select(folder)
    # Re-select source after dst operations (connection may have gone idle)
    src.select(src_q, readonly=True)

    for i in range(0, len(msg_ids), BATCH):
        chunk = msg_ids[i:i+BATCH]
        chunk_str = b','.join(chunk)
        try:
            st3, msg_data = src.fetch(chunk_str, '(INTERNALDATE FLAGS RFC822)')
        except Exception as e:
            errors += len(chunk)
            # Reconnect source
            try: src = connect_src(); src.select(src_q, readonly=True)
            except: pass
            continue

        for item in msg_data:
            if not isinstance(item, tuple): continue
            try:
                meta = item[0].decode('utf-8', errors='replace')
                raw_msg = item[1]

                # Stalwart v0.15.5 rejects \Flag syntax in APPEND — skip flags
                append_flags = None

                # Duplicate check by Message-ID
                parsed = email.message_from_bytes(raw_msg)
                mid = parsed.get('Message-ID', '').strip()
                if mid and mid in existing_ids:
                    skipped += 1; continue

                # Append
                dst.append(folder, append_flags, None, raw_msg)
                if mid: existing_ids.add(mid)
                copied += 1
            except Exception as e:
                errors += 1
                if errors <= 3:
                    sys.stdout.write(f'\n  ERR: {e}\n')
                    sys.stdout.flush()

        done = min(i+BATCH, src_count)
        sys.stdout.write(f'\r  {done}/{src_count} | ✓{copied} ⏭{skipped} ✗{errors}   ')
        sys.stdout.flush()

    print(f'\r  {src_count} src | ✓{copied} ⏭{skipped} ✗{errors}           ')
    return copied, skipped, errors

def main():
    print('=== Proton → Stalwart Migration ===')
    print(f'Started: {datetime.now():%Y-%m-%d %H:%M:%S}')
    src = connect_src(); dst = connect_dst()

    _, folders_raw = src.list()
    folders = []
    for f in folders_raw:
        d = f.decode()
        if 'Noselect' in d: continue
        parts = d.split('" "')
        if len(parts) < 2: continue
        folder = parts[-1].strip().strip('"')
        if folder in SKIP_FOLDERS: continue
        if any(folder.startswith(p) for p in SKIP_PREFIXES): continue
        folders.append(folder)

    print(f'Syncing {len(folders)} folders\n')
    results = []; t0 = time.time()
    total_c = total_s = total_e = 0

    for folder in folders:
        print(f'[{folder}]')
        try:
            c, s, e = sync_folder(src, dst, folder)
        except Exception as ex:
            print(f'  ✗ Fatal: {ex}')
            c, s, e = 0, 0, 1
            try: src = connect_src(); dst = connect_dst()
            except: pass
        results.append((folder, c, s, e))
        total_c += c; total_s += s; total_e += e

    elapsed = time.time() - t0
    print('\n' + '='*55)
    print('MIGRATION REPORT')
    print('='*55)
    print(f'Completed: {datetime.now():%Y-%m-%d %H:%M:%S}')
    print(f'Duration:  {int(elapsed//60)}m {int(elapsed%60)}s')
    print(f'Folders:   {len(folders)}')
    print(f'Copied:    {total_c}')
    print(f'Skipped:   {total_s} (duplicates)')
    print(f'Errors:    {total_e}')
    print()
    print(f'  {"Folder":<45} {"Copied":>7} {"Dup":>5} {"Err":>5}')
    print(f'  {"-"*45} {"-"*7} {"-"*5} {"-"*5}')
    for folder, c, s, e in results:
        mark = ' ⚠' if e > 0 else ''
        print(f'  {folder:<45} {c:>7} {s:>5} {e:>5}{mark}')
    src.logout(); dst.logout()

if __name__ == '__main__':
    main()