#!/usr/bin/env python3
"""Batch-update quota2 table for all active mailboxes from doveadm quota get."""

import subprocess
import sys

DOVEADM   = '/usr/local/dovecot/bin/doveadm'
MYSQL     = 'mysql'
DB_USER   = 'postfix'
DB_PASS   = 'T3CJnFMJNX9wmhNs'
DB_SOCKET = '/run/mysqld/mysqld.sock'
DB_NAME   = 'postfix'


def mysql_query(sql):
    return subprocess.run(
        [MYSQL, '-u', DB_USER, '-p' + DB_PASS, '--socket', DB_SOCKET,
         DB_NAME, '-sNe', sql],
        capture_output=True, text=True)


def get_users(limit=None):
    query = 'SELECT username FROM mailbox WHERE active = 1 ORDER BY username;'
    if limit:
        query = 'SELECT username FROM mailbox WHERE active = 1 ORDER BY username LIMIT %d;' % limit
    r = mysql_query(query)
    return [u.strip() for u in r.stdout.strip().split('\n') if u.strip()]


def get_quota(user):
    result = subprocess.run(
        [DOVEADM, 'quota', 'get', '-u', user],
        capture_output=True, text=True
    )
    bytes_val = None
    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
            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 = ("REPLACE INTO quota2 (username, bytes, messages) VALUES ('"
           + user.replace("'", "\'") + "', "
           + str(bytes_val) + ", " + str(messages_val) + ");")
    mysql_query(sql)


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)

    ok = 0
    skip = 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))
            skip += 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))
        ok += 1

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


if __name__ == '__main__':
    main()