#!/usr/bin/env python3
"""
Sync Keycloak users to Mailu mailboxes.
 - Generates/stores a mailu_app_password attribute in Keycloak (admin-only)
 - Upserts the mailbox in Mailu Postgres using that password
"""

import os
import sys
import json
import time
import secrets
import string
import datetime
import requests
import psycopg2
from psycopg2.extras import RealDictCursor
from passlib.hash import bcrypt_sha256


KC_BASE = os.environ["KEYCLOAK_BASE_URL"].rstrip("/")
KC_REALM = os.environ["KEYCLOAK_REALM"]
KC_CLIENT_ID = os.environ["KEYCLOAK_CLIENT_ID"]
KC_CLIENT_SECRET = os.environ["KEYCLOAK_CLIENT_SECRET"]

MAILU_DOMAIN = os.environ["MAILU_DOMAIN"]
MAILU_DEFAULT_QUOTA = int(os.environ.get("MAILU_DEFAULT_QUOTA", "20000000000"))

DB_CONFIG = {
    "host": os.environ["MAILU_DB_HOST"],
    "port": int(os.environ.get("MAILU_DB_PORT", "5432")),
    "dbname": os.environ["MAILU_DB_NAME"],
    "user": os.environ["MAILU_DB_USER"],
    "password": os.environ["MAILU_DB_PASSWORD"],
}

SESSION = requests.Session()


def log(msg):
    sys.stdout.write(f"{msg}\n")
    sys.stdout.flush()


def get_kc_token():
    resp = SESSION.post(
        f"{KC_BASE}/realms/{KC_REALM}/protocol/openid-connect/token",
        data={
            "grant_type": "client_credentials",
            "client_id": KC_CLIENT_ID,
            "client_secret": KC_CLIENT_SECRET,
        },
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()["access_token"]


def kc_get_users(token):
    users = []
    first = 0
    max_results = 200
    headers = {"Authorization": f"Bearer {token}"}
    while True:
        resp = SESSION.get(
            f"{KC_BASE}/admin/realms/{KC_REALM}/users",
            params={"first": first, "max": max_results, "enabled": "true"},
            headers=headers,
            timeout=20,
        )
        resp.raise_for_status()
        batch = resp.json()
        users.extend(batch)
        if len(batch) < max_results:
            break
        first += max_results
    return users


def kc_update_attributes(token, user, attributes):
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    payload = {
        "firstName": user.get("firstName"),
        "lastName": user.get("lastName"),
        "email": user.get("email"),
        "enabled": user.get("enabled", True),
        "username": user["username"],
        "emailVerified": user.get("emailVerified", False),
        "attributes": attributes,
    }
    user_url = f"{KC_BASE}/admin/realms/{KC_REALM}/users/{user['id']}"
    resp = SESSION.put(user_url, headers=headers, json=payload, timeout=20)
    resp.raise_for_status()
    verify = SESSION.get(
        user_url,
        headers={"Authorization": f"Bearer {token}"},
        params={"briefRepresentation": "false"},
        timeout=15,
    )
    verify.raise_for_status()
    attrs = verify.json().get("attributes") or {}
    if not attrs.get("mailu_app_password"):
        raise Exception(f"attribute not persisted for {user.get('email') or user['username']}")


def random_password():
    alphabet = string.ascii_letters + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(24))


def ensure_mailu_user(cursor, email, password, display_name):
    localpart, domain = email.split("@", 1)
    if domain.lower() != MAILU_DOMAIN.lower():
        return
    hashed = bcrypt_sha256.hash(password)
    now = datetime.datetime.utcnow()
    cursor.execute(
        """
        INSERT INTO "user" (
          email, localpart, domain_name, password,
          quota_bytes, quota_bytes_used,
          global_admin, enabled, enable_imap, enable_pop, allow_spoofing,
          forward_enabled, forward_destination, forward_keep,
          reply_enabled, reply_subject, reply_body, reply_startdate, reply_enddate,
          displayed_name, spam_enabled, spam_mark_as_read, spam_threshold,
          change_pw_next_login, created_at, updated_at, comment
        )
        VALUES (
          %(email)s, %(localpart)s, %(domain)s, %(password)s,
          %(quota)s, 0,
          false, true, true, true, false,
          false, '', true,
          false, NULL, NULL, DATE '1900-01-01', DATE '2999-12-31',
          %(display)s, true, true, 80,
          false, CURRENT_DATE, %(now)s, ''
        )
        ON CONFLICT (email) DO UPDATE
        SET password = EXCLUDED.password,
            enabled = true,
            updated_at = EXCLUDED.updated_at
        """,
        {
            "email": email,
            "localpart": localpart,
            "domain": domain,
            "password": hashed,
            "quota": MAILU_DEFAULT_QUOTA,
            "display": display_name or localpart,
            "now": now,
        },
    )


def main():
    token = get_kc_token()
    users = kc_get_users(token)
    if not users:
        log("No users found; exiting.")
        return

    conn = psycopg2.connect(**DB_CONFIG)
    conn.autocommit = True
    cursor = conn.cursor(cursor_factory=RealDictCursor)

    for user in users:
        attrs = user.get("attributes", {}) or {}
        app_pw_value = attrs.get("mailu_app_password")
        if isinstance(app_pw_value, list):
            app_pw = app_pw_value[0] if app_pw_value else None
        elif isinstance(app_pw_value, str):
            app_pw = app_pw_value
        else:
            app_pw = None

        email = user.get("email")
        if not email:
            email = f"{user['username']}@{MAILU_DOMAIN}"

        if not app_pw:
            app_pw = random_password()
            attrs["mailu_app_password"] = app_pw
            kc_update_attributes(token, user, attrs)
            log(f"Set mailu_app_password for {email}")

        display_name = " ".join(
            part for part in [user.get("firstName"), user.get("lastName")] if part
        ).strip()

        ensure_mailu_user(cursor, email, app_pw, display_name)
        log(f"Synced mailbox for {email}")

    cursor.close()
    conn.close()


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        log(f"ERROR: {exc}")
        sys.exit(1)