#!/usr/bin/env python3
"""
Batch-update the quota2 table for all active mailboxes.

Uses 'doveadm quota get' (must run as root) to read current usage and
writes it to quota2 via INSERT ... ON CONFLICT DO UPDATE.

Usage:
    update_quota2.sh [--dry-run] [--limit N]

Options:
    --dry-run   Print what would be done without touching the database.
    --limit N   Process only the first N mailboxes (useful for testing).
"""

import subprocess
import sys

DOVEADM = '/usr/local/dovecot/bin/doveadm'
PSQL    = 'psql'
DB_USER = 'postfix'
DB_NAME = 'postfix'


def get_users(limit=None):
    query = 'SELECT username FROM mailbox WHERE active = true ORDER BY username;'
    if limit:
        query = ('SELECT username FROM mailbox WHERE active = true '
                 'ORDER BY username LIMIT %d;' % limit)
    result = subprocess.run(
        [PSQL, '-U', DB_USER, DB_NAME, '-tAF', '|', '-c', query],
        capture_output=True, text=True
    )
    return [u.strip() for u in result.stdout.strip().split('\n') if u.strip()]


def get_quota(user):
    """Returns (bytes, messages) or (None, None) if not available."""
    result = subprocess.run(
        [DOVEADM, 'quota', 'get', '-u', user],
        capture_output=True, text=True
    )
    bytes_val = messages_val = None
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) < 4:
            continue
        if parts[2] == 'STORAGE' and parts[1] != 'Type':
            try:
                bytes_val = int(parts[3]) * 1024   # KiB → bytes
            except ValueError:
                bytes_val = 0
        elif parts[2] == 'MESSAGE' and parts[1] != 'Type':
            try:
                messages_val = int(parts[3])
            except ValueError:
                messages_val = 0
    return bytes_val, messages_val


def upsert_quota2(user, bytes_val, messages_val):
    sql = (
        "INSERT INTO quota2 (username, bytes, messages) VALUES ('"
        + user.replace("'", "''")
        + "', " + str(bytes_val) + ", " + str(messages_val) + ") "
        "ON CONFLICT (username) DO UPDATE "
        "SET bytes = EXCLUDED.bytes, messages = EXCLUDED.messages;"
    )
    subprocess.run([PSQL, '-U', DB_USER, DB_NAME, '-c', sql],
                   capture_output=True)


def main():
    dry_run = '--dry-run' in sys.argv
    limit = None
    if '--limit' in sys.argv:
        idx = sys.argv.index('--limit')
        try:
            limit = int(sys.argv[idx + 1])
        except (IndexError, ValueError):
            print('Usage: update_quota2.sh [--dry-run] [--limit N]')
            sys.exit(1)

    users = get_users(limit)
    total = len(users)
    print('Processing %d mailboxes...' % total)

    updated = skipped = 0
    for i, user in enumerate(users, 1):
        bytes_val, messages_val = get_quota(user)
        if bytes_val is None or messages_val is None:
            print('  [%d/%d] SKIP %s (no quota data)' % (i, total, user))
            skipped += 1
            continue
        if dry_run:
            print('  [%d/%d] DRY  %s: %d bytes, %d messages'
                  % (i, total, user, bytes_val, messages_val))
        else:
            upsert_quota2(user, bytes_val, messages_val)
            if i % 50 == 0 or i == total:
                print('  [%d/%d] done' % (i, total))
        updated += 1

    print('Done: %d updated, %d skipped.' % (updated, skipped))


if __name__ == '__main__':
    main()