#!/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:
                    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", "r") 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()