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
from pathlib import Path

import requests

DATA_DIR = Path(__file__).parent.parent.parent / "data"

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(DATA_DIR / "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()