ephemere/python/cohort/assign_cohort_role.py

#!/usr/bin/env python3
"""Assign the Cohort role to all 155 participants.
Respects Discord rate limits with proper backoff and retry logic.
"""

import json
import os
import time

import requests

BOT_TOKEN = os.environ["DISCORD_BOT_TOKEN"]
GUILD_ID = "692816967895220344"
COHORT_ROLE_ID = "1464314780935258112"

BASE_URL = "https://discord.com/api/v10"
HEADERS = {"Authorization": f"Bot {BOT_TOKEN}", "Content-Length": "0"}


def assign_role_with_retry(user_id: str, role_id: str, max_retries: int = 5) -> bool:
    url = f"{BASE_URL}/guilds/{GUILD_ID}/members/{user_id}/roles/{role_id}"

    for attempt in range(max_retries):
        response = requests.put(url, headers=HEADERS)

        if response.status_code == 204:
            return True
        elif response.status_code == 429:
            # Check headers first, fall back to JSON body
            retry_after = response.headers.get("Retry-After")
            if retry_after is None:
                retry_after = response.headers.get("X-RateLimit-Reset-After")
            if retry_after is None:
                try:
                    retry_after = response.json().get("retry_after", 1)
                except Exception:
                    retry_after = 1
            retry_after = float(retry_after)
            print(f"  Rate limited! Waiting {retry_after:.2f}s before retry...")
            time.sleep(retry_after)
        else:
            print(f"  Error {response.status_code}: {response.text}")
            backoff_time = (2**attempt) * 0.5
            print(f"  Retrying in {backoff_time:.2f}s...")
            time.sleep(backoff_time)

    return False


def main():
    with open("team_assignments.json") as f:
        teams = json.load(f)

    all_users = []
    for team in teams:
        all_users.extend(team["leaders"])
        all_users.extend(team["participants"])

    unique_users = list(dict.fromkeys(all_users))

    print(f"Assigning Cohort role to {len(unique_users)} users...")
    print(f"Role ID: {COHORT_ROLE_ID}")
    print("-" * 50)

    success_count = 0
    fail_count = 0

    for i, user_id in enumerate(unique_users, 1):
        print(f"[{i}/{len(unique_users)}] Assigning to {user_id}...", end=" ")

        if assign_role_with_retry(user_id, COHORT_ROLE_ID):
            print("✓")
            success_count += 1
        else:
            print("✗ FAILED")
            fail_count += 1

        # Small delay between requests to be nice to the API
        time.sleep(0.1)

    print("-" * 50)
    print(f"Complete! Success: {success_count}, Failed: {fail_count}")


if __name__ == "__main__":
    main()