modules/api_training_bridge.py

# api_training_bridge.py
# ================================================================
# PONTE DE TREINAMENTO VIA API - AKIRA IA V21
# ================================================================
# Conecta os provedores API (Cerebras, OpenRouter, Mistral, etc.)
# ao sistema de treinamento, permitindo que todos os 5 níveis
# usem APIs para distilar conhecimento.
# ================================================================

import os
import time
import json
from typing import Optional, List, Dict, Any, Tuple
from loguru import logger

try:
    from . import config
except (ImportError, SystemError):
    try:
        import modules.config as config
    except ImportError:
        config = None


class APITrainingBridge:
    """
    Ponte entre provedores API e sistema de treinamento.
    
    Cria clientes API leves para uso exclusivo de treinamento,
    reutilizando as mesmas chaves de API e rotações do LLMManager.
    
    Uso:
        bridge = APITrainingBridge()
        resposta, provider = bridge.generate("Analise a emoção desta mensagem: ...")
    """

    def __init__(self):
        self._clients = {}
        self._providers = []
        self._rotation = {}
        self._initialized = False
        self._stats = {}
        self._init_clients()

    def _init_clients(self):
        """Inicializa clientes API leves para treinamento."""
        if self._initialized:
            return

        # 1. Cerebras (OpenAI SDK compat)
        self._setup_cerebras()

        # 2. OpenRouter (OpenAI SDK compat)
        self._setup_openrouter()

        # 3. Mistral (requests direto)
        self._setup_mistral()

        # 4. Groq (Groq SDK)
        self._setup_groq()

        # 5. Gemini (Google GenAI)
        self._setup_gemini()

        # 6. Cohere
        self._setup_cohere()

        # 7. HF Inference
        self._setup_hf_inference()

        self._initialized = True

        if self._providers:
            logger.info(
                f"🔗 [API-TRAINING] Bridge inicializado: "
                f"{', '.join(self._providers)} ({len(self._providers)} providers)"
            )
        else:
            logger.warning("🔗 [API-TRAINING] Nenhum provider disponível para treinamento")

    def _setup_cerebras(self):
        """Cerebras via rotação multi-conta."""
        try:
            from .cerebras_rotation import get_cerebras_rotation
            rotation = get_cerebras_rotation()
            if rotation.account_names:
                import openai
                key = rotation.get_current_api_key()
                if key:
                    client = openai.OpenAI(
                        api_key=key,
                        base_url="https://api.cerebras.ai/v1",
                        timeout=30.0,
                        max_retries=0,
                    )
                    self._clients["cerebras"] = client
                    self._rotation["cerebras"] = rotation
                    self._providers.append("cerebras")
                    logger.debug(f"✅ [API-TRAINING] Cerebras OK ({rotation.get_current_account_name()})")
        except Exception as e:
            logger.debug(f"[API-TRAINING] Cerebras skip: {e}")

    def _setup_openrouter(self):
        """OpenRouter via rotação multi-conta."""
        try:
            from .openrouter_rotation import get_openrouter_rotation
            rotation = get_openrouter_rotation()
            if rotation.api_keys:
                import openai
                key = rotation.get_current_key()
                if key:
                    client = openai.OpenAI(
                        api_key=key,
                        base_url="https://openrouter.ai/api/v1",
                        timeout=30.0,
                        max_retries=0,
                    )
                    self._clients["openrouter"] = client
                    self._rotation["openrouter"] = rotation
                    self._providers.append("openrouter")
                    logger.debug(f"✅ [API-TRAINING] OpenRouter OK ({rotation.get_current_account_name()})")
        except Exception as e:
            logger.debug(f"[API-TRAINING] OpenRouter skip: {e}")

    def _setup_mistral(self):
        """Mistral via API key direta ou rotação."""
        try:
            from .mistral_rotation import get_mistral_rotation
            rotation = get_mistral_rotation(config)
            if rotation:
                self._rotation["mistral"] = rotation
                self._providers.append("mistral")
                logger.debug(f"✅ [API-TRAINING] Mistral OK ({rotation.get_current_account_name()})")
            elif config and getattr(config, "MISTRAL_API_KEY", ""):
                self._providers.append("mistral")
                logger.debug("✅ [API-TRAINING] Mistral OK (chave única)")
        except Exception as e:
            logger.debug(f"[API-TRAINING] Mistral skip: {e}")

    def _setup_groq(self):
        """Groq SDK."""
        try:
            api_key = getattr(config, "GROQ_API_KEY", "") if config else ""
            if api_key and len(api_key) > 5:
                from groq import Groq
                client = Groq(api_key=api_key)
                self._clients["groq"] = client
                self._providers.append("groq")
                logger.debug("✅ [API-TRAINING] Groq OK")
        except Exception as e:
            logger.debug(f"[API-TRAINING] Groq skip: {e}")

    def _setup_gemini(self):
        """Google Gemini via genai."""
        try:
            if config:
                gemini_key = getattr(config, "GEMINI_API_KEY", "") or getattr(config, "GOOGLE_API_KEY", "")
                if gemini_key and len(gemini_key) > 5:
                    import google.generativeai as genai
                    genai.configure(api_key=gemini_key)
                    model_name = getattr(config, "GEMINI_MODEL", "gemini-2.5-flash")
                    self._clients["gemini"] = genai.GenerativeModel(model_name)
                    self._providers.append("gemini")
                    logger.debug("✅ [API-TRAINING] Gemini OK")
        except Exception as e:
            logger.debug(f"[API-TRAINING] Gemini skip: {e}")

    def _setup_cohere(self):
        """Cohere SDK."""
        try:
            api_key = getattr(config, "COHERE_API_KEY", "") if config else ""
            if api_key and len(api_key) > 5:
                from cohere import Client
                client = Client(api_key=api_key)
                self._clients["cohere"] = client
                self._providers.append("cohere")
                logger.debug("✅ [API-TRAINING] Cohere OK")
        except Exception as e:
            logger.debug(f"[API-TRAINING] Cohere skip: {e}")

    def _setup_hf_inference(self):
        """HF Inference via rotação multi-conta."""
        try:
            from .hf_inference_rotation import get_hf_inference_rotation
            rotation = get_hf_inference_rotation()
            configured = [
                acc for acc in rotation.account_order
                if os.getenv(rotation.accounts[acc])
            ]
            if configured:
                from huggingface_hub import InferenceClient
                token = rotation.get_current_api_token()
                if token:
                    client = InferenceClient(token=token, timeout=30.0)
                    self._clients["hf_inference"] = client
                    self._rotation["hf_inference"] = rotation
                    self._providers.append("hf_inference")
                    logger.debug(f"✅ [API-TRAINING] HF Inference OK ({len(configured)} contas)")
        except Exception as e:
            logger.debug(f"[API-TRAINING] HF Inference skip: {e}")

    # ============================================================
    # INTERFACE PÚBLICA
    # ============================================================

    def generate(
        self,
        prompt: str,
        system: str = "Responda de forma clara e direta.",
        max_tokens: int = 2048,
        temperature: float = 0.7,
        preferred_provider: Optional[str] = None,
    ) -> Tuple[str, str]:
        """
        Gera resposta via API para treinamento.
        
        Returns:
            (resposta_texto, provider_name) ou ("", "") se todos falharem
        """
        providers = self._providers[:]
        if preferred_provider and preferred_provider in providers:
            providers.remove(preferred_provider)
            providers.insert(0, preferred_provider)

        for provider in providers:
            try:
                text = self._call_provider(provider, prompt, system, max_tokens, temperature)
                if text and len(text.strip()) > 5:
                    self._record_call(provider, True)
                    return text.strip(), provider
                self._record_call(provider, False)
            except Exception as e:
                self._record_call(provider, False)
                err_str = str(e)
                logger.debug(f"[API-TRAINING] {provider} falhou: {e}")
                # 429 (rate limit temporário) → rotaciona conta
                if "429" in err_str or "rate" in err_str.lower():
                    self._handle_rate_limit(provider)
                # 402 (quota/crédito esgotado) → já tratado em _call_cerebras;
                # para outros providers, tenta rotacionar também.
                elif "402" in err_str or "payment" in err_str.lower():
                    self._handle_rate_limit(provider)

        return "", ""

    def generate_batch(
        self,
        prompts: List[Dict[str, str]],
        max_concurrent: int = 1,
        delay_between: float = 1.0,
    ) -> List[Dict[str, Any]]:
        """
        Gera respostas para batch de prompts.
        
        Args:
            prompts: Lista de {"prompt": str, "system": str (opcional)}
            max_concurrent: Máximo de chamadas simultâneas
            delay_between: Delay entre chamadas (rate limiting)
        
        Returns:
            Lista de {"prompt": str, "response": str, "provider": str, "success": bool}
        """
        results = []
        for i, item in enumerate(prompts):
            prompt = item.get("prompt", "")
            system = item.get("system", "Responda de forma clara e direta.")
            
            text, provider = self.generate(prompt, system)
            
            results.append({
                "prompt": prompt,
                "response": text,
                "provider": provider,
                "success": bool(text),
            })
            
            if delay_between > 0 and i < len(prompts) - 1:
                time.sleep(delay_between)
        
        return results

    def get_best_provider(self) -> str:
        """Retorna o provider com melhor taxa de sucesso."""
        best = ""
        best_score = -1
        
        for provider, stats in self._stats.items():
            total = stats.get("success", 0) + stats.get("failure", 0)
            if total < 3:
                continue
            rate = stats["success"] / total
            if rate > best_score:
                best_score = rate
                best = provider
        
        return best or (self._providers[0] if self._providers else "")

    def get_stats(self) -> Dict[str, Any]:
        """Retorna estatísticas de uso."""
        return {
            "providers_available": self._providers[:],
            "stats": dict(self._stats),
        }

    def is_available(self) -> bool:
        """Retorna se há providers disponíveis."""
        return bool(self._providers)

    def get_mac_training_context(self) -> Dict[str, Any]:
        """
        Retorna contexto de treinamento com estado do drive MAC.

        Inclui informações de saúde, uso e performance do disco
        para contextualizar dados de treinamento em sistemas Apple.
        """
        import subprocess

        context = {
            "platform": "unknown",
            "drives": [],
            "timestamp": time.time(),
        }

        try:
            import platform as plat
            system = plat.system()

            if system == "Darwin":
                context["platform"] = "macos"
                context["drives"] = self._get_mac_drive_states()
            else:
                context["platform"] = system.lower()
                context["drives"] = self._get_generic_drive_states()
        except Exception as e:
            logger.debug(f"[API-TRAINING] MAC context error: {e}")

        return context

    def _get_mac_drive_states(self) -> List[Dict[str, Any]]:
        """Coleta estado dos drives em macOS via diskutil."""
        drives = []
        try:
            result = subprocess.run(
                ["diskutil", "list", "-plist"],
                capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0 and result.stdout:
                plist_data = __import__("plistlib").loads(result.stdout.encode())
                for disk in plist_data.get("AllDisks", []):
                    info = plist_data.get("AllDiskInformation", {}).get(disk, {})
                    drives.append({
                        "name": disk,
                        "volume_name": info.get("VolumeName", ""),
                        "mount_point": info.get("MountPoint", ""),
                        "total_bytes": info.get("TotalSize", 0),
                        "free_bytes": info.get("FreeSpace", 0),
                        "filesystem": info.get("Filesystem", ""),
                        "protocol": info.get("BusProtocol", ""),
                    })
        except Exception as e:
            logger.debug(f"[API-TRAINING] diskutil error: {e}")

        return drives

    def _get_generic_drive_states(self) -> List[Dict[str, Any]]:
        """Coleta estado dos drives em sistemas genéricos (Windows/Linux)."""
        drives = []
        try:
            import shutil
            usage = shutil.disk_usage("/")
            drives.append({
                "name": "/",
                "total_bytes": usage.total,
                "free_bytes": usage.free,
                "used_bytes": usage.used,
            })
        except Exception as e:
            logger.debug(f"[API-TRAINING] disk usage error: {e}")

        return drives

    # ============================================================
    # CHAMADAS POR PROVIDER
    # ============================================================

    def _call_provider(
        self, provider: str, prompt: str, system: str,
        max_tokens: int, temperature: float
    ) -> Optional[str]:
        """Roteia chamada para o provider correto."""
        if provider == "cerebras":
            return self._call_cerebras(prompt, system, max_tokens, temperature)
        elif provider == "openrouter":
            return self._call_openrouter(prompt, system, max_tokens, temperature)
        elif provider == "mistral":
            return self._call_mistral(prompt, system, max_tokens, temperature)
        elif provider == "groq":
            return self._call_groq(prompt, system, max_tokens, temperature)
        elif provider == "gemini":
            return self._call_gemini(prompt, system, max_tokens, temperature)
        elif provider == "cohere":
            return self._call_cohere(prompt, system, max_tokens, temperature)
        elif provider == "hf_inference":
            return self._call_hf_inference(prompt, system, max_tokens, temperature)
        return None

    def _call_cerebras(self, prompt, system, max_tokens, temperature):
        client = self._clients.get("cerebras")
        if not client:
            return None
        # Modelo Cerebras serverless default: gpt-oss-120b (120B params, 65k contexto, production).
        # Sobrescreve via config.CEREBRAS_MODEL ou env CEREBRAS_MODEL.
        model = getattr(config, "CEREBRAS_MODEL", "gpt-oss-120b") if config else "gpt-oss-120b"
        try:
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return response.choices[0].message.content if response.choices else None
        except Exception as e:
            err_str = str(e)
            # 402 Payment Required → quota esgotada (não é rate limit temporário).
            # Blacklist 24h da conta atual e rotaciona, igual ao api.py.
            if "402" in err_str or "payment_required" in err_str.lower():
                rotation = self._rotation.get("cerebras")
                if rotation and hasattr(rotation, "handle_quota_402"):
                    try:
                        current_name = rotation.get_current_account_name()
                        rotation.handle_quota_402(current_name)
                        # Atualiza o cliente para a próxima conta disponível
                        new_key = rotation.get_current_api_key()
                        if new_key and new_key != client.api_key:
                            import openai
                            self._clients["cerebras"] = openai.OpenAI(
                                api_key=new_key,
                                base_url="https://api.cerebras.ai/v1",
                                timeout=30.0,
                                max_retries=0,
                            )
                            logger.info(f"🔄 [API-TRAINING] Cerebras 402 → rotacionado para '{rotation.get_current_account_name()}'")
                    except Exception as rotate_e:
                        logger.debug(f"[API-TRAINING] Cerebras 402 rotate error: {rotate_e}")
            raise

    def _call_openrouter(self, prompt, system, max_tokens, temperature, rotation=None):
        rotation = rotation or self._rotation.get("openrouter")
        client = self._clients.get("openrouter")
        if not client:
            return None

        # Atualiza chave se tem rotação
        if rotation:
            key = rotation.get_current_key()
            if key:
                client.api_key = key

        model = getattr(config, "OPENROUTER_MODEL", "nvidia/nemotron-3-ultra-550b-a55b:free") if config else "nvidia/nemotron-3-ultra-550b-a55b:free"
        try:
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": system},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return response.choices[0].message.content if response.choices else None
        except Exception as e:
            err_str = str(e)
            # 429 → rotaciona conta OpenRouter
            if "429" in err_str or "rate" in err_str.lower():
                if rotation and hasattr(rotation, "handle_429_error"):
                    try:
                        rotation.handle_429_error()
                        new_key = rotation.get_current_key()
                        if new_key:
                            client.api_key = new_key
                    except Exception:
                        pass
            raise

    def _call_mistral(self, prompt, system, max_tokens, temperature):
        import requests
        
        rotation = self._rotation.get("mistral")
        api_key = ""
        if rotation:
            api_key = rotation.get_current_key()
        elif config:
            api_key = getattr(config, "MISTRAL_API_KEY", "")
        
        if not api_key:
            return None
        
        model = getattr(config, "MISTRAL_MODEL", "mistral-large-latest") if config else "mistral-large-latest"
        response = requests.post(
            "https://api.mistral.ai/v1/chat/completions",
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json={
                "model": model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": prompt},
                ],
                "max_tokens": max_tokens,
                "temperature": temperature,
            },
            timeout=30,
        )
        if response.status_code == 200:
            data = response.json()
            choices = data.get("choices", [])
            return choices[0]["message"]["content"] if choices else None
        elif response.status_code == 429:
            # Rate limit → rotaciona conta Mistral e levanta para o generate() tratar
            rotation = self._rotation.get("mistral")
            if rotation and hasattr(rotation, "handle_429_error"):
                try:
                    rotation.handle_429_error()
                except Exception:
                    pass
            raise Exception("429 Rate Limit")
        elif response.status_code == 402:
            # Payment Required → quota esgotada
            raise Exception("402 Payment Required")
        return None

    def _call_groq(self, prompt, system, max_tokens, temperature):
        client = self._clients.get("groq")
        if not client:
            return None
        if len(system or '') + len(prompt or '') > 7500:
            return None
        model = getattr(config, "GROQ_MODEL", "groq/compound") if config else "groq/compound"
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ],
            max_tokens=max_tokens,
            temperature=temperature,
        )
        return response.choices[0].message.content if response.choices else None

    def _call_gemini(self, prompt, system, max_tokens, temperature):
        model = self._clients.get("gemini")
        if not model:
            return None
        full_prompt = f"{system}\n\n{prompt}"
        response = model.generate_content(
            full_prompt,
            generation_config={
                "max_output_tokens": max_tokens,
                "temperature": temperature,
            },
        )
        return response.text if response.text else None

    def _call_cohere(self, prompt, system, max_tokens, temperature):
        client = self._clients.get("cohere")
        if not client:
            return None
        model = getattr(config, "COHERE_MODEL", "command-r-plus-08-2024") if config else "command-r-plus-08-2024"
        response = client.chat(
            model=model,
            message=prompt,
            preamble=system,
            max_tokens=max_tokens,
            temperature=temperature,
        )
        return response.text if response.text else None

    def _call_hf_inference(self, prompt, system, max_tokens, temperature):
        client = self._clients.get("hf_inference")
        if not client:
            return None
        rotation = self._rotation.get("hf_inference")
        model = getattr(config, "MISTRAL_MODEL_HF", "mistralai/Mistral-7B-Instruct-v0.2") if config else "mistralai/Mistral-7B-Instruct-v0.2"
        
        # Atualiza token se tem rotação
        if rotation:
            token = rotation.get_current_api_token()
            if token:
                client.token = token
        
        full_prompt = f"{system}\n\n{prompt}"
        response = client.text_generation(
            full_prompt,
            model=model,
            max_new_tokens=max_tokens,
            temperature=temperature,
        )
        return response if isinstance(response, str) else None

    # ============================================================
    # UTILITÁRIOS
    # ============================================================

    def _record_call(self, provider: str, success: bool):
        """Registra resultado de chamada."""
        if provider not in self._stats:
            self._stats[provider] = {"success": 0, "failure": 0}
        if success:
            self._stats[provider]["success"] += 1
        else:
            self._stats[provider]["failure"] += 1

    def _handle_rate_limit(self, provider: str):
        """Trata rate limit rotacionando para próxima conta."""
        rotation = self._rotation.get(provider)
        if not rotation:
            return
        # Diferentes rotações usam nomes diferentes de método para 429.
        # Cerebras: handle_rate_limit_error(); OpenRouter/Mistral: handle_429_error().
        handler = getattr(rotation, "handle_429_error", None) or getattr(rotation, "handle_rate_limit_error", None)
        if handler:
            try:
                handler()
                logger.info(f"🔄 [API-TRAINING] {provider} rotacionado após rate limit")
            except Exception:
                pass


# Singleton
_bridge_instance = None

def get_api_training_bridge() -> APITrainingBridge:
    """Obtém ou cria singleton do bridge."""
    global _bridge_instance
    if _bridge_instance is None:
        _bridge_instance = APITrainingBridge()
    return _bridge_instance