import redis
import json
import random

class KeyManager:
    def __init__(self):
        self.r = redis.Redis(
            host='10.99.0.3', 
            port=6379, 
            password='clube67_dragonfly_pass_9903', 
            db=0,
            decode_responses=True
        )
        self.pool_key = 'swarm:keys:pool'

    def get_key(self, provider):
        """Busca uma chave ativa para o provedor (google ou anthropic)"""
        pool_data = self.r.get(self.pool_key)
        if not pool_data:
            return None
        
        pool = json.loads(pool_data)
        keys = pool.get(provider, [])
        
        # Filtra apenas as chaves ativas
        active_keys = [k for k in keys if k.get('status') == 'active']
        
        if not active_keys:
            print(f" [!] Erro: Nenhuma chave ativa para {provider}")
            return None
        
        # Seleciona uma chave aleatória do pool ativo (Load Balancing)
        selected = random.choice(active_keys)
        return selected['key']

    def report_failure(self, provider, failed_key):
        """Marca uma chave como inativa caso falhe (ex: sem saldo)"""
        pool_data = self.r.get(self.pool_key)
        pool = json.loads(pool_data)
        
        for k in pool.get(provider, []):
            if k['key'] == failed_key:
                k['status'] = 'inactive'
                print(f" [W] Chave {failed_key[:10]}... marcada como INATIVA.")
        
        self.r.set(self.pool_key, json.dumps(pool))

# Exemplo de uso:
# km = KeyManager()
# key = km.get_key('google')