#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Script standalone pour vérifier et afficher les propriétés d'un token Hugging Face.
Ce script peut être exécuté séparément pour diagnostiquer les problèmes d'authentification.
"""

import os
import sys
import json
import requests
from datetime import datetime
from dotenv import load_dotenv
from argparse import ArgumentParser


def color_text(text, color_code):
    """Format text with color for terminal output."""
    return f"\033[{color_code}m{text}\033[0m"


def success(text):
    """Format text as success message (green)."""
    return color_text(f"✅ {text}", "92")


def warning(text):
    """Format text as warning message (yellow)."""
    return color_text(f"⚠️ {text}", "93")


def error(text):
    """Format text as error message (red)."""
    return color_text(f"❌ {text}", "91")


def info(text):
    """Format text as info message (blue)."""
    return color_text(f"ℹ️ {text}", "94")


def check_token_via_inference_api(token=None, verbose=True):
    """
    Vérifie la validité d'un token HF en testant directement l'API d'inférence.
    L'API whoami ne fonctionne pas toujours correctement pour les tokens mais l'API d'inférence
    est la priorité dans notre application.
    
    Args:
        token: Le token à vérifier
        verbose: Afficher des informations détaillées
    
    Returns:
        dict: Résultats de la vérification
    """
    results = {
        "is_valid": False,
        "token": None,
        "error_message": None,
        "can_access_inference": False
    }
    
    # 1. Obtenir le token
    if token is None:
        token = os.environ.get("HF_TOKEN")
    
    if not token:
        print(error("Aucun token trouvé. Veuillez spécifier un token avec --token ou définir la variable d'environnement HF_TOKEN."))
        results["error_message"] = "No token provided"
        return results
    
    # Ne montrer aucun caractère du token, juste indiquer sa présence
    masked_token = "••••••••••"
    results["token"] = masked_token
    
    print(info(f"Token à vérifier: {masked_token}"))
    
    # 2. Vérifier le format basique
    if not token.startswith("hf_"):
        print(warning("Le token ne commence pas par 'hf_' ce qui est inhabituel. Vérifiez son format."))
    else:
        print(success("Format du token valide (commence par 'hf_')"))
    
    # 3. Tester l'API d'inférence directement - méthode recommandée pour valider un token
    try:
        # Test avec un modèle public simple
        test_model = "gpt2"
        api_url = f"https://api-inference.huggingface.co/models/{test_model}"
        
        print(info(f"Test du token avec l'API d'inférence sur le modèle public {test_model}..."))
        
        headers = {"Authorization": f"Bearer {token}"}
        payload = {"inputs": "Hello, how are you?"}
        
        response = requests.post(api_url, headers=headers, json=payload, timeout=10)
        
        if response.status_code in [200, 503]:  # 503 signifie que le modèle est en cours de chargement, mais le token est valide
            print(success(f"Token valide pour l'API d'inférence! Status code: {response.status_code}"))
            if response.status_code == 503:
                print(info("Le modèle est en cours de chargement. Le token a bien été accepté par l'API."))
            results["is_valid"] = True
            results["can_access_inference"] = True
            
            if verbose and response.status_code == 200:
                print(info("Résultat de l'inférence:"))
                print(json.dumps(response.json(), indent=2))
        else:
            print(error(f"Échec du test de l'API d'inférence. Status code: {response.status_code}"))
            results["error_message"] = response.text
            
            try:
                error_data = response.json()
                if "error" in error_data:
                    print(error(f"Message d'erreur: {error_data['error']}"))
                    results["error_message"] = error_data['error']
            except:
                print(error(f"Message d'erreur: {response.text}"))
            
            # En cas d'échec, tester aussi l'endpoint de liste des modèles
            try:
                print(info("Test alternatif avec la liste des modèles déployés..."))
                list_url = "https://api-inference.huggingface.co/status"
                list_response = requests.get(list_url, headers=headers, timeout=10)
                
                if list_response.status_code == 200:
                    print(success("Le token peut accéder à la liste des modèles déployés"))
                    results["can_access_inference"] = True
                    results["is_valid"] = True
                else:
                    print(error(f"Échec de l'accès à la liste des modèles. Status code: {list_response.status_code}"))
            except Exception as e:
                print(error(f"Erreur lors du test alternatif: {str(e)}"))
                
    except Exception as e:
        print(error(f"Erreur lors du test de l'API d'inférence: {str(e)}"))
        results["error_message"] = str(e)
    
    # 4. Tests supplémentaires des permissions
    if results["is_valid"]:
        try:
            print(info("\nTest des permissions du token..."))
            
            # Tester si on peut accéder aux modèles privés de l'organisation
            if os.environ.get("HF_ORGANIZATION"):
                org = os.environ.get("HF_ORGANIZATION")
                print(info(f"Test d'accès aux modèles de l'organisation {org}..."))
                
                # On regarde juste si on peut accéder à la liste des modèles de l'organisation
                org_url = f"https://huggingface.co/api/models?author={org}"
                org_response = requests.get(org_url, headers=headers, timeout=10)
                
                if org_response.status_code == 200:
                    print(success(f"Accès autorisé aux modèles de l'organisation {org}"))
                else:
                    print(warning(f"Le token n'a pas accès aux modèles de l'organisation {org}"))
        except Exception as e:
            print(error(f"Erreur lors du test des permissions: {str(e)}"))
    
    return results


def check_model_access(token, model, verbose=False):
    """
    Vérifie si le token a accès à un modèle spécifique.
    
    Args:
        token: Token HF à vérifier
        model: Nom du modèle à tester
        verbose: Afficher des informations détaillées
    
    Returns:
        bool: True si le modèle est accessible, False sinon
    """
    print(f"\n" + info(f"Test d'accès au modèle: {model}"))
    
    headers = {
        "Authorization": f"Bearer {token}"
    }
    
    # 1. Vérifier si le modèle existe et est accessible via l'API d'inférence
    try:
        api_url = f"https://api-inference.huggingface.co/models/{model}"
        payload = {"inputs": "Hello, test access"}
        
        print(info(f"Test d'accès à l'API d'inférence pour {model}..."))
        
        response = requests.post(api_url, headers=headers, json=payload, timeout=20)
        
        if response.status_code in [200, 503]:  # 503 = modèle en cours de chargement, mais le token est valide
            if response.status_code == 200:
                print(success(f"Accès réussi à l'API d'inférence pour {model}"))
                return True
            else:
                print(success(f"Accès autorisé pour {model} (modèle en cours de chargement)"))
                return True
        else:
            error_message = "Unknown error"
            try:
                error_data = response.json()
                if "error" in error_data:
                    error_message = error_data["error"]
            except:
                error_message = response.text
                
            print(error(f"Échec d'accès à l'API d'inférence pour {model}: {response.status_code}"))
            print(error(f"Message: {error_message}"))
            
            # Analyse de l'erreur
            if "quota" in error_message.lower() or "rate" in error_message.lower():
                print(warning("Possible problème de quota ou de limite de taux"))
            elif "loading" in error_message.lower():
                print(info("Le modèle est en cours de chargement - réessayez plus tard"))
                return True  # Considérer comme un succès car le token est accepté
            elif "permission" in error_message.lower() or "access" in error_message.lower():
                print(error("Problème de permissions - vous n'avez pas accès à ce modèle"))
            
            # Faire un test alternatif via l'API du Hub
            try:
                print(info(f"Test alternatif via l'API du Hub pour {model}..."))
                hub_url = f"https://huggingface.co/api/models/{model}"
                hub_response = requests.get(hub_url, headers=headers, timeout=10)
                
                if hub_response.status_code == 200:
                    print(warning(f"Le modèle {model} existe et est accessible via l'API Hub, mais pas via l'API d'inférence"))
                    print(info("Cela peut être dû à des restrictions sur le modèle ou à des problèmes temporaires de l'API"))
                    if verbose:
                        model_info = hub_response.json()
                        if model_info.get("private", False):
                            print(info("Ce modèle est privé"))
                        if model_info.get("gated", False):
                            print(info("Ce modèle est à accès restreint (gated)"))
                else:
                    print(error(f"Le modèle {model} n'est pas accessible via l'API Hub non plus: {hub_response.status_code}"))
            except Exception as e:
                print(error(f"Erreur lors du test alternatif: {str(e)}"))
                
            return False
            
    except Exception as e:
        print(error(f"Erreur lors du test d'accès au modèle: {str(e)}"))
        return False


def main():
    parser = ArgumentParser(description="Vérifiez les propriétés d'un token Hugging Face")
    parser.add_argument("--token", type=str, help="Token Hugging Face à vérifier (si non spécifié, utilise HF_TOKEN)")
    parser.add_argument("--verbose", "-v", action="store_true", help="Afficher des informations détaillées")
    parser.add_argument("--test-model", "-m", type=str, help="Tester l'accès à un modèle spécifique")
    parser.add_argument("--test-premium", action="store_true", help="Tester l'accès aux modèles premium courants")
    
    args = parser.parse_args()
    
    # Charger les variables d'environnement
    load_dotenv()
    
    print(info(f"=== Vérification de Token Hugging Face - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ===\n"))
    
    # Vérifier le token via l'API d'inférence directement
    token = args.token or os.environ.get("HF_TOKEN")
    token_info = check_token_via_inference_api(token, args.verbose)
    
    # Si le token est valide et qu'on a demandé de tester un modèle
    if token_info["is_valid"]:
        if args.test_model:
            check_model_access(token, args.test_model, args.verbose)
        
        if args.test_premium:
            print("\n" + info("=== Test d'accès aux modèles premium ==="))
            premium_models = [
                "meta-llama/Llama-3.3-70B-Instruct",
                "mistralai/Mistral-Small-24B-Instruct-2501",
                "deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
            ]
            
            for model in premium_models:
                result = check_model_access(token, model, args.verbose)
                print(info(f"Résultat pour {model}: {success('Accessible') if result else error('Non accessible')}"))
                print("-" * 50)


if __name__ == "__main__":
    main()