# Authentik expression policy: sync a member's Discord ROLES into Authentik GROUPS at every login.
# Paste into: Customization -> Policies -> Create -> Expression Policy  (name: discord-roles-to-groups)
# Bind to BOTH flows (after discord-guild-member): default-source-enrollment and default-source-authentication.
# Tick "Execution logging" on the policy: everything below shows up in Events -> Logs -> policy_execution
# under result.log_messages, and in `docker logs authentik-server`.
# Requires the "guilds.members.read" scope on the Discord source (source -> Protocol settings -> Scopes:
#   guilds guilds.members.read ). Existing users re-consent once at Discord.
# The Authentik groups must already exist (Directory -> Groups). Role IDs: Discord Developer Mode ->
#   Server Settings -> Roles -> right-click a role -> Copy Role ID.
# Manual memberships are safe: the policy only removes groups it granted itself (user attribute discord_role_groups).

GUILD_ID = "1387203044084551690"                        # <- your Discord server ID

# Discord role id : (label used in the log, [Authentik groups this role grants])
# One role may grant several groups, and several roles may grant the same group. Do NOT repeat a role id:
# this is a dict, so a second entry for the same id silently replaces the first.
ROLES = {
    "1387239222070612110": ("Discord Admins", ["bot-users", "admins"]),
    "1515901670943490129": ("Lead Dev",       ["bot-users"]),
    "1566623319027880086": ("Dev",            ["bot-users"]),
#   "222222222222222222": ("Streamer",        ["streamers"]),
#   "333333333333333333": ("Editor",          ["editors"]),
#   "444444444444444444": ("Media admin",     ["media-admins"]),
#   "555555555555555555": ("Git user",        ["git-users"]),
}

from authentik.core.models import Group

source = context.get("source")
if not source or getattr(source, "provider_type", "") != "discord":
    return True                                          # not a Discord login (e.g. akadmin): nothing to sync

connection = context.get("goauthentik.io/sources/connection")
if not connection:
    ak_logger.info("discord-roles: no source connection in context; nothing to sync")
    return True

resp = requests.get(
    f"https://discord.com/api/users/@me/guilds/{GUILD_ID}/member",
    headers={"Authorization": f"Bearer {connection.access_token}"},
    timeout=10,
)
if resp.status_code != 200:
    ak_logger.info(f"discord-roles: member lookup failed, HTTP {resp.status_code}: {resp.text[:200]}")
    ak_message(f"Could not read your Discord roles ({resp.status_code}); groups unchanged.")
    return True                                          # don't block login on a Discord hiccup

roles = set(resp.json().get("roles", []))

# ---- what Discord says, what we map, what we ignore ------------------------------------------------
matched = {rid: ROLES[rid] for rid in roles if rid in ROLES}
unmapped = sorted(roles - set(ROLES))
wanted = {group for _, groups in matched.values() for group in groups}
managed = {group for _, groups in ROLES.values() for group in groups}

ak_logger.info(f"discord-roles: discord returned {len(roles)} role(s): {sorted(roles)}")
for rid, (label, groups) in sorted(matched.items(), key=lambda item: item[1][0]):
    ak_logger.info(f"discord-roles: MAPPED   role {rid} ({label}) -> {groups}")
if unmapped:
    ak_logger.info(f"discord-roles: ignored  {len(unmapped)} role(s) not in the table: {unmapped}")
if not matched:
    ak_logger.info(f"discord-roles: no mapped roles; the table knows {sorted(ROLES)}")
ak_logger.info(f"discord-roles: groups wanted {sorted(wanted)}; managed by this policy {sorted(managed)}; "
               f"withheld {sorted(managed - wanted)}")

# A group named in the table but absent from Authentik can never be granted: say so instead of failing quietly.
existing = {group.name: group for group in Group.objects.filter(name__in=managed)}
absent = sorted(managed - set(existing))
if absent:
    ak_logger.info(f"discord-roles: WARNING these groups do not exist in Authentik and were skipped: {absent}")

# ---- apply ------------------------------------------------------------------------------------------
# Authentication flow: the user exists -> edit memberships directly.
# Only memberships THIS policy granted earlier are ever removed (tracked in user.attributes["discord_role_groups"]),
# so a group someone was added to by hand in Authentik survives role changes and re-logins.
user = context.get("pending_user") or (request.user if request.user.is_authenticated else None)
if user and user.pk:
    attrs = user.attributes or {}
    granted = set(attrs.get("discord_role_groups", []))
    added, removed = [], []
    for name, group in existing.items():
        if name in wanted and name not in granted:
            group.users.add(user); granted.add(name); added.append(name)
        elif name in wanted:
            group.users.add(user)                 # idempotent: make sure the membership is really there
        elif name in granted:
            group.users.remove(user)                  # role gone -> take back what the role had given
            granted.discard(name); removed.append(name)
    attrs["discord_role_groups"] = sorted(granted)
    user.attributes = attrs
    user.save()
    ak_logger.info(f"discord-roles: authentication flow, user {user.username} (pk {user.pk}): "
                   f"added {added or '[]'}, removed {removed or '[]'}, role-granted now {sorted(granted)}")
    return True

# Enrollment flow: the user is created later by the user-write stage; hand it the groups to attach and
# record them as role-granted so a later role removal can revoke exactly these. The flow plan is reached
# through context["flow_plan"] on some versions and IS context on others, so write to whichever is there.
attach = [group for name, group in existing.items() if name in wanted]
plan = context.get("flow_plan")
target = plan.context if plan is not None else context
target["groups"] = attach
prompt = target.setdefault("prompt_data", {})
prompt.setdefault("attributes", {})["discord_role_groups"] = sorted(wanted)
ak_logger.info(f"discord-roles: enrollment flow via {'flow_plan' if plan is not None else 'policy context'}: "
               f"attaching {[group.name for group in attach]}, recording {sorted(wanted)}")
return True