#!/usr/bin/env python3
"""
Set or remove the 'admin' custom claim on a Firebase user.

Usage (run from drb-c2-core directory):
  python scripts/set_admin.py grant user@example.com
  python scripts/set_admin.py revoke user@example.com

Requires GCP_CREDENTIALS_PATH or Application Default Credentials.
The user must sign out and back in (or wait up to 1 hour) for the
new claim to take effect in their ID token.
"""
import sys
import os
import firebase_admin
from firebase_admin import credentials, auth

def main():
    if len(sys.argv) != 3 or sys.argv[1] not in ("grant", "revoke"):
        print(__doc__)
        sys.exit(1)

    action, email = sys.argv[1], sys.argv[2]

    creds_path = os.getenv("GCP_CREDENTIALS_PATH", "gcp-key.json")
    cred = credentials.Certificate(creds_path)
    firebase_admin.initialize_app(cred)

    try:
        user = auth.get_user_by_email(email)
    except auth.UserNotFoundError:
        print(f"No Firebase user found for {email!r}")
        sys.exit(1)

    existing = user.custom_claims or {}
    if action == "grant":
        updated = {**existing, "admin": True}
        auth.set_custom_user_claims(user.uid, updated)
        print(f"Admin granted to {email} ({user.uid})")
    else:
        updated = {k: v for k, v in existing.items() if k != "admin"}
        auth.set_custom_user_claims(user.uid, updated)
        print(f"Admin revoked from {email} ({user.uid})")

    print("The user must sign out and back in for the change to take effect.")

if __name__ == "__main__":
    main()