#!/usr/bin/env python3
"""
assign-roles.py — Attribue les rôles RBAC aux users dans Zitadel.

Pour chaque user actif Furious :
  1. Détermine son rôle (via mapping job_title)
  2. Cherche son compte Zitadel par email
  3. Crée le UserGrant avec le rôle dans le project "Infra Seize"

Prérequis :
  - 88 users déjà importés dans Zitadel (cf. import-furious-to-zitadel.py)
  - 8 rôles déjà créés dans le project Zitadel "Infra Seize"
  - Clé service account avec rôle IAM_USER_MANAGER

Auteur : Hedi Kalboussi
Date : 5 juin 2026
"""

import json
import logging
import sys
import time
from pathlib import Path

import jwt
import requests
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend

# Réutiliser le même mapping
from generate_roles_mapping import JOB_TITLE_TO_ROLE, determine_role  # noqa
# Note : adapter selon disposition réelle des fichiers
# Si import échoue, copier le mapping ici directement

# ============================================================================
# CONFIGURATION
# ============================================================================

SCRIPT_DIR = Path(__file__).parent
KEY_FILE = SCRIPT_DIR / "zitadelfuriousimportkey.json"
RAW_DATA_FILE = SCRIPT_DIR / "furious-users-raw.json"

ZITADEL_DOMAIN = "https://sso.seizeconsulting.com"
PROJECT_ID = "375123185571463171"  # ⚠️ ID du project "Infra Seize" — adapter si différent

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

# ============================================================================
# AUTHENTIFICATION ZITADEL (identique au script d'import)
# ============================================================================

def get_zitadel_token():
    with open(KEY_FILE) as f:
        key_data = json.load(f)
    user_id = key_data["userId"]
    key_id = key_data["keyId"]
    private_key_pem = key_data["key"]
    private_key = serialization.load_pem_private_key(
        private_key_pem.encode(), password=None, backend=default_backend()
    )
    now = int(time.time())
    payload = {"iss": user_id, "sub": user_id, "aud": ZITADEL_DOMAIN, "iat": now, "exp": now + 3600}
    assertion = jwt.encode(payload, private_key, algorithm="RS256", headers={"kid": key_id})
    response = requests.post(
        f"{ZITADEL_DOMAIN}/oauth/v2/token",
        data={
            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
            "scope": "openid profile urn:zitadel:iam:org:project:id:zitadel:aud",
            "assertion": assertion,
        },
    )
    response.raise_for_status()
    return response.json()["access_token"]

# ============================================================================
# API ZITADEL — Search user by email + create user grant
# ============================================================================

def find_user_by_email(email, token):
    """Cherche un user Zitadel par email. Retourne l'userId."""
    response = requests.post(
        f"{ZITADEL_DOMAIN}/v2/users",
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
        json={
            "queries": [
                {"emailQuery": {"emailAddress": email, "method": "TEXT_QUERY_METHOD_EQUALS"}}
            ]
        },
    )
    if not response.ok:
        log.error(f"Erreur recherche {email} : {response.status_code}")
        return None
    
    result = response.json()
    users = result.get("result", [])
    if not users:
        return None
    return users[0].get("userId")


def assign_role_to_user(user_id, role_key, token):
    """Attribue un rôle au user dans le project Infra Seize."""
    response = requests.post(
        f"{ZITADEL_DOMAIN}/management/v1/users/{user_id}/grants",
        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
        json={
            "projectId": PROJECT_ID,
            "roleKeys": [role_key],
        },
    )
    
    if response.status_code == 409:
        log.info(f"  Grant déjà existant pour user {user_id}")
        return True
    
    if not response.ok:
        log.error(f"Échec attribution rôle pour {user_id} : {response.status_code} {response.text[:200]}")
        return False
    
    return True

# ============================================================================
# MAIN
# ============================================================================

def main():
    log.info("===== Attribution des rôles RBAC =====")
    
    with open(RAW_DATA_FILE) as f:
        data = json.load(f)
    
    actifs = [
        u for u in data.get("data", {}).get("User", [])
        if u.get("departure_date") == "0000-00-00"
    ]
    log.info(f"Users actifs : {len(actifs)}")
    
    token = get_zitadel_token()
    log.info("✅ Token Zitadel OK")
    
    success_count = 0
    no_role_count = 0
    not_found_count = 0
    error_count = 0
    
    for i, user in enumerate(actifs, 1):
        email = user.get("email", "").lower()
        role = determine_role(user)
        
        log.info(f"[{i}/{len(actifs)}] {email} → rôle: {role}")
        
        if not role:
            log.warning(f"  Pas de rôle déterminé (job_title='{user.get('job_title')}')")
            no_role_count += 1
            continue
        
        user_id = find_user_by_email(email, token)
        if not user_id:
            log.error(f"  User Zitadel introuvable pour {email}")
            not_found_count += 1
            continue
        
        if assign_role_to_user(user_id, role, token):
            success_count += 1
        else:
            error_count += 1
        
        time.sleep(0.15)  # Anti rate-limit
    
    log.info(f"===== FIN =====")
    log.info(f"  ✅ Succès        : {success_count}")
    log.info(f"  ⚠️ Sans rôle     : {no_role_count}")
    log.info(f"  ❌ Non trouvés   : {not_found_count}")
    log.info(f"  💥 Erreurs       : {error_count}")


if __name__ == "__main__":
    main()