#!/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"))
MAILU_ENABLED_ATTR = os.environ.get("MAILU_ENABLED_ATTR", "mailu_enabled")
MAILU_EMAIL_ATTR = "mailu_email"
MAILU_SYSTEM_USERS = [
    item.strip()
    for item in os.environ.get("MAILU_SYSTEM_USERS", "").split(",")
    if item.strip()
]
MAILU_SYSTEM_PASSWORD = os.environ.get("MAILU_SYSTEM_PASSWORD", "").strip()

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 retry_request(label, func, attempts=10):
    for attempt in range(1, attempts + 1):
        try:
            return func()
        except requests.RequestException as exc:
            if attempt == attempts:
                raise
            log(f"{label} failed (attempt {attempt}/{attempts}): {exc}")
            time.sleep(attempt * 2)


def retry_db_connect(attempts=10):
    for attempt in range(1, attempts + 1):
        try:
            return psycopg2.connect(**DB_CONFIG)
        except psycopg2.Error as exc:
            if attempt == attempts:
                raise
            log(f"Database connection failed (attempt {attempt}/{attempts}): {exc}")
            time.sleep(attempt * 2)


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",
                "briefRepresentation": "false",
            },
            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",
    }
    user_url = f"{KC_BASE}/admin/realms/{KC_REALM}/users/{user['id']}"
    current = SESSION.get(
        user_url,
        headers={"Authorization": f"Bearer {token}"},
        params={"briefRepresentation": "false"},
        timeout=15,
    )
    current.raise_for_status()
    current_payload = current.json()
    current_attrs = current_payload.get("attributes") if isinstance(current_payload, dict) else None
    if not isinstance(current_attrs, dict):
        current_attrs = {}
    current_attrs.update(attributes)
    resp = SESSION.put(user_url, headers=headers, json={"attributes": current_attrs}, 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 get_attribute_value(attributes, key):
    raw = (attributes or {}).get(key)
    if isinstance(raw, list):
        return raw[0] if raw else None
    if isinstance(raw, str):
        return raw
    return None


def mailu_enabled(attributes) -> bool:
    raw = get_attribute_value(attributes, MAILU_ENABLED_ATTR)
    if raw is None:
        return bool(get_attribute_value(attributes, MAILU_EMAIL_ATTR))
    return str(raw).strip().lower() in {"1", "true", "yes", "y", "on"}


def resolve_mailu_email(user, attributes):
    explicit = get_attribute_value(attributes, MAILU_EMAIL_ATTR)
    if explicit:
        return explicit

    email = user.get("email") or ""
    if "@" in email and email.lower().endswith(f"@{MAILU_DOMAIN.lower()}"):
        return email

    return f"{user['username']}@{MAILU_DOMAIN}"


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.now(datetime.timezone.utc)
    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 ensure_system_mailboxes(cursor):
    if not MAILU_SYSTEM_USERS:
        return
    if not MAILU_SYSTEM_PASSWORD:
        log("MAILU_SYSTEM_USERS set but MAILU_SYSTEM_PASSWORD is missing; skipping system mailboxes")
        return

    for email in MAILU_SYSTEM_USERS:
        localpart = email.split("@", 1)[0] if "@" in email else email
        try:
            ensure_mailu_user(cursor, email, MAILU_SYSTEM_PASSWORD, localpart)
            log(f"Ensured system mailbox for {email}")
        except Exception as exc:
            log(f"Failed to ensure system mailbox {email}: {exc}")


def main():
    token = retry_request("Keycloak token", get_kc_token)
    users = retry_request("Keycloak user list", lambda: kc_get_users(token))
    if not users and not MAILU_SYSTEM_USERS:
        log("No users found; exiting.")
        return

    conn = retry_db_connect()
    conn.autocommit = True
    cursor = conn.cursor(cursor_factory=RealDictCursor)

    for user in users:
        attrs = user.get("attributes", {}) or {}
        if user.get("enabled") is False:
            continue
        needs_update = False
        if get_attribute_value(attrs, MAILU_ENABLED_ATTR) is None and get_attribute_value(attrs, MAILU_EMAIL_ATTR):
            attrs[MAILU_ENABLED_ATTR] = ["true"]
            needs_update = True
        if not mailu_enabled(attrs):
            continue
        app_pw = get_attribute_value(attrs, "mailu_app_password")
        mailu_email = resolve_mailu_email(user, attrs)

        if not get_attribute_value(attrs, MAILU_EMAIL_ATTR):
            attrs[MAILU_EMAIL_ATTR] = [mailu_email]
            needs_update = True

        if not app_pw:
            app_pw = random_password()
            attrs["mailu_app_password"] = [app_pw]
            needs_update = True

        if needs_update:
            kc_update_attributes(token, user, attrs)
            log(f"Updated Mailu attributes for {mailu_email}")

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

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

    ensure_system_mailboxes(cursor)

    cursor.close()
    conn.close()


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