modules/cloudflare_rotation.py

"""
Cloudflare Workers AI Rotation Manager
Rotação multi-conta para Cloudflare Workers AI (SDXL, FLUX).
Free tier: 10,000 requests/dia por conta.

Cada conta precisa de ACCOUNT_ID + API_TOKEN.
Contas descobertas automaticamente via env vars:

  Primária:
    CLOUDFLARE_ACCOUNT_ID + CLOUDFLARE_API_TOKEN

  Extra (sufixo _2, _3, ...):
    CLOUDFLARE_ACCOUNT_ID_2 + CLOUDFLARE_API_TOKEN_2
    CLOUDFLARE_ACCOUNT_ID_3 + CLOUDFLARE_API_TOKEN_3
    ... até _9
"""

import os
import logging
from datetime import datetime, timedelta
from typing import Dict, Optional, List, Tuple

logger = logging.getLogger(__name__)


class CloudflareRotation:
    _instance: Optional['CloudflareRotation'] = None

    def __init__(self):
        # {account_name: (account_id_env, api_token_env)}
        self.accounts: Dict[str, Tuple[str, str]] = {}
        self._discover_accounts()

        self.current_account_idx = 0
        self.account_order = list(self.accounts.keys())
        self.rate_limited: Dict[str, tuple] = {}

        self._log_status()

    @staticmethod
    def _resolve_token_env() -> str:
        """Retorna o nome da env var do token, tentando variantes."""
        candidates = [
            "CLOUDFLARE_API_TOKEN",
            "CLOUDFLARE_API_KEY",
            "CLOUDFALRE_API_KEY",
        ]
        for c in candidates:
            if os.getenv(c):
                return c
        return "CLOUDFLARE_API_TOKEN"

    def _discover_accounts(self):
        """Descobre contas: primária + sufixos _2, _3, etc."""
        token_env = self._resolve_token_env()
        self.accounts['main'] = ('CLOUDFLARE_ACCOUNT_ID', token_env)

        suffix = 2
        while suffix <= 9:
            acc_id_env = f'CLOUDFLARE_ACCOUNT_ID_{suffix}'
            for base in ("CLOUDFLARE_API_TOKEN", "CLOUDFLARE_API_KEY", "CLOUDFALRE_API_KEY"):
                variant = f"{base}_{suffix}"
                if os.getenv(acc_id_env) and os.getenv(variant):
                    self.accounts[f'main_{suffix}'] = (acc_id_env, variant)
                    break
            else:
                break
            suffix += 1

        # Filtra apenas contas completamente configuradas
        self.accounts = {
            name: envs for name, envs in self.accounts.items()
            if os.getenv(envs[0]) and os.getenv(envs[1])
        }

    def _log_status(self):
        if self.account_order:
            logger.info(
                f"🌐 [CLOUDFLARE ROTATION] {len(self.account_order)} conta(s): "
                f"{', '.join(self.account_order)}"
            )
        else:
            logger.warning(
                "⚠️ [CLOUDFLARE ROTATION] Nenhuma conta configurada. "
                "Define CLOUDFLARE_ACCOUNT_ID + CLOUDFLARE_API_TOKEN no .env"
            )

    def get_current_account_name(self) -> str:
        if not self.account_order:
            return "nenhuma"
        return self.account_order[self.current_account_idx]

    def get_current_credentials(self) -> Optional[Tuple[str, str]]:
        if not self.account_order:
            return None
        attempts = 0
        max_attempts = len(self.account_order)

        while attempts < max_attempts:
            current_account = self.get_current_account_name()
            if not self.is_account_limited(current_account):
                acc_id_env, token_env = self.accounts[current_account]
                acc_id = os.getenv(acc_id_env)
                api_token = os.getenv(token_env)
                if acc_id and api_token:
                    logger.debug(f"?? [CLOUDFLARE] Usando conta: {current_account}")
                    return (acc_id, api_token)
                # Sem credenciais: rodar + contar (senão o ciclo era infinito
                # e pendurava quem chamava no arranque).
                logger.debug(f"?? [CLOUDFLARE] Conta {current_account} sem credenciais, a rodar...")
                self.rotate_to_next()
                attempts += 1
            else:
                logger.debug(f"⏭️ [CLOUDFLARE] Conta {current_account} limitada, rotacionando...")
                self.rotate_to_next()
                attempts += 1

        logger.error("❌ [CLOUDFLARE] Todas as contas limitadas ou mal configuradas")
        return None

    def rotate_to_next(self) -> str:
        if not self.account_order:
            return "nenhuma"
        self.current_account_idx = (self.current_account_idx + 1) % len(self.account_order)
        new_account = self.get_current_account_name()
        logger.info(f"🔄 [CLOUDFLARE ROTATION] Switched to: {new_account}")
        return new_account

    def handle_rate_limit_error(self, error_msg: str = None) -> str:
        current_account = self.get_current_account_name()
        limited_until = datetime.now() + timedelta(minutes=10)
        self.rate_limited[current_account] = (limited_until, error_msg or "Rate limit")
        logger.warning(f"⚠️ [CLOUDFLARE] Conta {current_account} limitada. Erro: {error_msg}")
        return self.rotate_to_next()

    def is_account_limited(self, account_name: str) -> bool:
        if account_name not in self.rate_limited:
            return False
        limited_until, _ = self.rate_limited[account_name]
        if datetime.now() < limited_until:
            return True
        del self.rate_limited[account_name]
        return False

    def get_all_credentials(self) -> List[Tuple[str, str]]:
        creds = []
        for acc_name in self.account_order:
            if not self.is_account_limited(acc_name):
                acc_id_env, token_env = self.accounts[acc_name]
                acc_id = os.getenv(acc_id_env)
                api_token = os.getenv(token_env)
                if acc_id and api_token:
                    creds.append((acc_id, api_token))
        return creds


def get_cloudflare_rotation() -> CloudflareRotation:
    if CloudflareRotation._instance is None:
        CloudflareRotation._instance = CloudflareRotation()
    return CloudflareRotation._instance