modules/context_builder.py

# type: ignore
"""
================================================================================
AKIRA V21 ULTIMATE - CONTEXT BUILDER MODULE
================================================================================
Constrói prompts otimizados para LLM combinando:
- Memória de curto prazo (100 mensagens)
- Contexto de reply (prioritário)
- Memória vetorial (fatos aprendidos)
- Contexto emocional
- Sistema adaptativo baseado em tamanho da pergunta

Features:
- Hierarquia correta de contexto (reply > curto prazo > vetorial)
- Token budgeting inteligente
- Ajuste adaptativo para perguntas curtas
- Suporte a múltiplos provedores LLM
================================================================================
"""

import os
import sys
import time
import json
import logging
from typing import Optional, Dict, Any, List, Tuple
from dataclasses import dataclass

# Imports robustos com fallback - CORRIGIDO para usar modules.
try:
    from . import config
    from .context_isolation import ContextIsolationManager, ConversationContext
    from .short_term_memory import ShortTermMemory, MessageWithContext
    from .reply_context_handler import ReplyContextHandler, ProcessedReplyContext
    CONTEXT_BUILDER_AVAILABLE = True
except ImportError:
    try:
        import modules.config as config
        from modules.context_isolation import ContextIsolationManager, ConversationContext
        from modules.short_term_memory import ShortTermMemory, MessageWithContext
        from modules.reply_context_handler import ReplyContextHandler, ProcessedReplyContext
        CONTEXT_BUILDER_AVAILABLE = True
    except ImportError:
        CONTEXT_BUILDER_AVAILABLE = False
        config = None

logger = logging.getLogger(__name__)

# ============================================================
# CONFIGURAÇÃO
# ============================================================

# Token budgets para diferentes componentes
TOKEN_BUDGET_SYSTEM: int = 1500
TOKEN_BUDGET_REPLY: int = 800  # Para contexto de reply
TOKEN_BUDGET_SHORT_TERM: int = 4000  # Para memória de curto prazo
TOKEN_BUDGET_VECTOR: int = 1000  # Para memória vetorial
TOKEN_BUDGET_TOTAL: int = 8000  # Total disponível para contexto

# Limiares para perguntas curtas
SHORT_QUESTION_THRESHOLD: int = 5  # palavras


@dataclass
class PromptBuildResult:
    """
    Resultado da construção do prompt.
    
    Attributes:
        system_prompt: Prompt do sistema (sem modificação)
        full_prompt: Prompt completo com contexto
        context_sections: Seções de contexto incluídas
        token_counts: Contagem de tokens por seção
        warnings: Avisos sobre limitações
        should_use_vector_memory: Se deve usar memória vetorial
        should_prioritize_reply: Se reply deve ser priorizado
    """
    system_prompt: str = ""
    full_prompt: str = ""
    context_sections: Dict[str, str] = None
    token_counts: Dict[str, int] = None
    warnings: List[str] = None
    should_use_vector_memory: bool = True
    should_prioritize_reply: bool = False
    
    def __post_init__(self):
        if self.context_sections is None:
            self.context_sections = {}
        if self.token_counts is None:
            self.token_counts = {}
        if self.warnings is None:
            self.warnings = []


# ============================================================
# FUNÇÕES AUXILIARES
# ============================================================

def estimar_tokens(texto: str) -> int:
    """Estima tokens em um texto (aproximação para português)."""
    if not texto:
        return 0
    # Média de 4 caracteres por token em português
    return max(1, len(texto) // 4)


def truncar_para_tokens(texto: str, max_tokens: int) -> str:
    """Trunca texto para caber no limite de tokens."""
    if not texto or max_tokens <= 0:
        return ""
    
    tokens = texto.split()
    if len(tokens) <= max_tokens:
        return texto
    
    return " ".join(tokens[:max_tokens])


def is_pergunta_curta(texto: str) -> bool:
    """Verifica se é uma pergunta curta."""
    if not texto:
        return False
    return len(texto.split()) <= SHORT_QUESTION_THRESHOLD


def calcular_peso_contexto(
    mensagem: str,
    reply_context: Optional[ProcessedReplyContext] = None
) -> float:
    """
    Calcula peso do contexto baseado no tamanho da mensagem e reply.
    
    Args:
        mensagem: Mensagem do usuário
        reply_context: Contexto de reply (opcional)
        
    Returns:
        Float entre 0.5 e 1.0 representando peso do contexto geral
    """
    word_count = len(mensagem.split())
    
    # Pergunta muito curta = menos contexto geral necessário
    if word_count <= 2:
        return 0.5
    
    # Pergunta curta = contexto moderado
    if word_count <= SHORT_QUESTION_THRESHOLD:
        return 0.7
    
    # Pergunta normal = contexto completo
    return 1.0


# ============================================================
# CLASSE PRINCIPAL
# ============================================================

class ContextBuilder:
    """
    Construtor de prompts otimizados para LLM.
    
    Hierarquia de contexto:
    1. System prompt (fixo)
    2. Reply context (prioritário se existir)
    3. Short-term memory (100 msgs sliding window)
    4. Vector memory (fatos aprendidos)
    5. User message (última)
    
    Adaptação para perguntas curtas:
    - Pergunta curta + reply: reply tem 100%, contexto geral 50%
    - Pergunta curta sem reply: contexto geral 70%
    - Pergunta normal: contexto geral 100%
    """
    
    def __init__(self, config_module=None):
        """
        Inicializa o builder.
        
        Args:
            config_module: Módulo de configuração (usa config se None)
        """
        self.config = config_module or config
        self.isolation_manager = None
        self._initialized = False
        
        if CONTEXT_BUILDER_AVAILABLE:
            try:
                self.isolation_manager = ContextIsolationManager()
                self._initialized = True
            except Exception as e:
                logger.warning(f"ContextBuilder: falha ao init isolation: {e}")
    
    def _ensure_initialized(self):
        """Garante inicialização."""
        if not self._initialized and CONTEXT_BUILDER_AVAILABLE:
            try:
                self.isolation_manager = ContextIsolationManager()
                self._initialized = True
            except:
                pass
    
    def build_prompt(
        self,
        user_message: str,
        conversation_id: str,
        system_prompt: str = None,
        reply_context: Optional[ProcessedReplyContext] = None,
        short_term_memory: Optional[ShortTermMemory] = None,
        vector_memory_info: Optional[List[Dict[str, Any]]] = None,
        emocao_atual: str = "neutral",
        incluir_memoria_vetorial: bool = True,
        max_tokens_contexto: int = TOKEN_BUDGET_TOTAL
    ) -> PromptBuildResult:
        """
        Constrói prompt completo para LLM.
        
        Args:
            user_message: Mensagem do usuário
            conversation_id: ID da conversa isolada
            system_prompt: Prompt do sistema (usa config se None)
            reply_context: Contexto de reply (opcional)
            short_term_memory: Memória de curto prazo (opcional)
            vector_memory_info: Fatos da memória vetorial (opcional)
            emocao_atual: Emoção atual do usuário
            incluir_memoria_vetorial: Se deve incluir memória vetorial
            max_tokens_contexto: Máximo de tokens para contexto
            
        Returns:
            PromptBuildResult com prompt completo
        """
        result = PromptBuildResult()
        
        # Get system prompt
        system_prompt = system_prompt or getattr(self.config, 'SYSTEM_PROMPT', '')
        result.system_prompt = system_prompt
        
        # Inicializa seções
        sections = {
            "system": system_prompt,
            "reply_context": "",
            "short_term_context": "",
            "vector_memory": "",
            "emotional_context": "",
            "user_message": user_message
        }
        
        # Contadores de tokens
        tokens = {
            "system": estimar_tokens(system_prompt),
            "reply": 0,
            "short_term": 0,
            "vector": 0,
            "emotional": 0,
            "user": estimar_tokens(user_message)
        }
        
        # Remaining budget after system and user
        remaining_budget = max_tokens_contexto - tokens["system"] - tokens["user"]
        
        # ===== 1. REPLY CONTEXT (PRIORITÁRIO!) =====
        if reply_context and reply_context.is_reply:
            result.should_prioritize_reply = True
            
            # Para perguntas curtas com reply, mais tokens para reply
            if is_pergunta_curta(user_message):
                reply_budget = min(TOKEN_BUDGET_REPLY * 1.5, int(remaining_budget * 0.35))
                remaining_budget -= reply_budget
            else:
                reply_budget = min(TOKEN_BUDGET_REPLY, int(remaining_budget * 0.25))
                remaining_budget -= reply_budget
            
            # Constrói section do reply
            reply_section = self._build_reply_section(reply_context, user_message)
            reply_section = truncar_para_tokens(reply_section, reply_budget)
            
            sections["reply_context"] = reply_section
            tokens["reply"] = estimar_tokens(reply_section)
        
        # ===== 2. SHORT-TERM MEMORY =====
        if short_term_memory:
            # Calcula peso baseado em tamanho da pergunta
            peso_contexto = calcular_peso_contexto(user_message, reply_context)
            stm_budget = min(
                int(TOKEN_BUDGET_SHORT_TERM * peso_contexto),
                int(remaining_budget * 0.7)
            )
            
            stm_section = self._build_short_term_section(
                short_term_memory, 
                reply_context,
                stm_budget
            )
            
            sections["short_term_context"] = stm_section
            tokens["short_term"] = estimar_tokens(stm_section)
            remaining_budget -= tokens["short_term"]
        
        # ===== 3. VECTOR MEMORY =====
        if incluir_memoria_vetorial and vector_memory_info:
            vector_budget = min(TOKEN_BUDGET_VECTOR, int(remaining_budget * 0.3))
            
            vector_section = self._build_vector_section(vector_memory_info, vector_budget)
            
            sections["vector_memory"] = vector_section
            tokens["vector"] = estimar_tokens(vector_section)
            remaining_budget -= tokens["vector"]
        
        # ===== 4. EMOTIONAL CONTEXT =====
        emotional_section = self._build_emotional_section(emocao_atual)
        sections["emotional_context"] = emotional_section
        tokens["emotional"] = estimar_tokens(emotional_section)
        
        # ===== 5. MONTA PROMPT COMPLETO =====
        prompt_parts = []
        
        # System
        if sections["system"]:
            prompt_parts.append(f"[SYSTEM]\n{sections['system']}\n[/SYSTEM]\n")
        
        # Emotional context (apenas se não neutral)
        if sections["emotional_context"]:
            prompt_parts.append(f"[EMOÇÃO ATUAL]\n{sections['emotional_context']}\n")
        
        # Reply context (prioritário!)
        if sections["reply_context"]:
            prompt_parts.append(f"[REPLY PRIORITÁRIO]\n{sections['reply_context']}\n")
        
        # Short-term context
        if sections["short_term_context"]:
            prompt_parts.append(f"[CONTEXTO RECENTE]\n{sections['short_term_context']}\n")
        
        # Vector memory
        if sections["vector_memory"]:
            prompt_parts.append(f"[MEMÓRIA APRENDIDA]\n{sections['vector_memory']}\n")
        
        # User message
        prompt_parts.append(f"[MENSAGEM]\n{user_message}\n")
        
        result.full_prompt = "\n".join(prompt_parts)
        result.context_sections = sections
        result.token_counts = tokens
        
        # Warnings se orçamento estourado
        total_tokens = sum(tokens.values())
        if total_tokens > max_tokens_contexto:
            result.warnings.append(f"Contexto grande: {total_tokens} tokens (limite: {max_tokens_contexto})")
        
        return result
    
    def _build_reply_section(
        self,
        reply_context: ProcessedReplyContext,
        user_message: str
    ) -> str:
        """Constrói seção de reply priorizado."""
        parts = []
        
        # Cabeçalho de prioridade
        if reply_context.priority_level >= 4:  # CRÍTICO
            parts.append("⚠️⚠️⚠️ REPLY CRÍTICO - PERGUNTA CURTA ⚠️⚠️⚠️")
        elif reply_context.priority_level == 3:  # REPLY TO BOT
            parts.append("⚠️ REPLY DIRETO AO BOT")
        else:
            parts.append("📎 REPLY")
        
        # Autor
        if reply_context.reply_to_bot:
            parts.append("Você está sendo diretamente mencionado!")
        else:
            parts.append(f"Respondendo a: {reply_context.quoted_author_name}")
        
        # Mensagem citada
        if reply_context.mensagem_citada:
            cited = reply_context.mensagem_citada[:300]
            parts.append(f"\nMsg citada:\n{cited}")
        
        # Contexto hint
        if reply_context.context_hint and reply_context.context_hint != "contexto_geral":
            parts.append(f"\nContexto: {reply_context.context_hint}")
        
        return "\n".join(parts)
    
    def _build_short_term_section(
        self,
        short_term_memory: ShortTermMemory,
        reply_context: Optional[ProcessedReplyContext] = None,
        max_tokens: int = TOKEN_BUDGET_SHORT_TERM
    ) -> str:
        """Constrói seção de memória de curto prazo."""
        # Obtém mensagens do contexto
        messages = short_term_memory.get_context_window(
            include_replies=True,
            prioritize_replies=True,
            max_tokens=max_tokens
        )
        
        if not messages:
            return ""
        
        parts = []
        parts.append("(últimas mensagens - replies priorizados)")
        
        # Limita a quantidade para caber no orçamento
        included_count = 0
        current_tokens = 0
        
        for msg in messages:
            msg_tokens = estimar_tokens(msg.content)
            if current_tokens + msg_tokens > max_tokens:
                break
            
            # Formata mensagem
            role = "🤖" if msg.role == "assistant" else "👤"
            content_preview = msg.content[:100] + ("..." if len(msg.content) > 100 else "")
            
            if msg.is_reply:
                parts.append(f"{role} [REPLY] {content_preview}")
            else:
                parts.append(f"{role} {content_preview}")
            
            current_tokens += msg_tokens
            included_count += 1
        
        if not parts:
            return ""
        
        return "\n".join(parts)
    
    def _build_vector_section(
        self,
        vector_info: List[Dict[str, Any]],
        max_tokens: int = TOKEN_BUDGET_VECTOR
    ) -> str:
        """Constrói seção de memória vetorial."""
        if not vector_info:
            return ""
        
        parts = []
        parts.append("(fatos aprendidos nesta conversa)")
        
        current_tokens = 0
        
        for item in vector_info[:10]:  # Limita a 10 itens
            text = item.get("text", "") or item.get("mensagem", "")
            if not text:
                continue
            
            text_preview = text[:80] + ("..." if len(text) > 80 else "")
            current_tokens += estimar_tokens(text)
            
            if current_tokens > max_tokens:
                break
            
            parts.append(f"• {text_preview}")
        
        if len(parts) == 1:
            return ""
        
        return "\n".join(parts)
    
    def _build_emotional_section(self, emocao: str) -> str:
        """Constrói seção de contexto emocional."""
        if emocao in ["neutral", "neutro"]:
            return ""
        
        emocoes_descritas = {
            "joy": "usuário parece feliz/contento",
            "felicidade": "usuário parece feliz/contento",
            "tristeza": "usuário parece triste",
            "triste": "usuário parece triste",
            "raiva": "usuário parece irritado/raivoso",
            "raivoso": "usuário parece irritado/raivoso",
            "amor": "usuário demonstra afeto",
            "medo": "usuário parece preocupado/assustado",
            "surpresa": "usuário parece surpreso",
            "surpreso": "usuário parece surpreso"
        }
        
        descricao = emocoes_descritas.get(emocao.lower(), f"usuário parece {emocao}")
        return f"Tom emocional: {descricao}"
    
    # ============================================================
    # HELPERS PARA API
    # ============================================================
    
    def build_history_for_llm(
        self,
        short_term_memory: ShortTermMemory,
        reply_context: Optional[ProcessedReplyContext] = None,
        max_tokens: int = TOKEN_BUDGET_SHORT_TERM
    ) -> List[Dict[str, str]]:
        """
        Constrói histórico formatado para LLM.
        
        Args:
            short_term_memory: Memória de curto prazo
            reply_context: Contexto de reply (opcional)
            max_tokens: Máximo de tokens
            
        Returns:
            Lista de dicts com role e content
        """
        # Garante que reply_context está priorizado
        if reply_context and reply_context.is_reply:
            # Cria mensagem artificial para o reply
            reply_entry = {
                "role": "user",
                "content": f"[REPLY] {reply_context.get_reply_summary_for_llm(reply_context)}"
            }
            
            # Obtém resto do histórico
            history = short_term_memory.get_messages_for_llm(
                reply_context=None,  # Já adicionado
                max_tokens=max_tokens - estimar_tokens(reply_entry["content"])
            )
            
            # Insere reply no início
            return [reply_entry] + history
        
        return short_term_memory.get_messages_for_llm(max_tokens=max_tokens)
    
    def estimate_prompt_tokens(
        self,
        user_message: str,
        reply_context: Optional[ProcessedReplyContext] = None,
        historico_size: int = 0
    ) -> int:
        """
        Estima tokens totais do prompt.
        
        Args:
            user_message: Mensagem do usuário
            reply_context: Contexto de reply
            historico_size: Tamanho do histórico em mensagens
            
        Returns:
            Estimativa de tokens
        """
        system_tokens = TOKEN_BUDGET_SYSTEM
        
        reply_tokens = 0
        if reply_context and reply_context.is_reply:
            reply_tokens = TOKEN_BUDGET_REPLY
        
        history_tokens = historico_size * 50  # Aproximação
        
        return system_tokens + reply_tokens + history_tokens + estimar_tokens(user_message)
    
    def get_conversation_context(
        self,
        numero_usuario: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None
    ) -> Tuple[Optional[ConversationContext], ShortTermMemory]:
        """
        Obtém contexto isolado e memória de curto prazo.
        
        Args:
            numero_usuario: Número do usuário
            tipo_conversa: "pv" ou "grupo"
            grupo_id: ID do grupo
            
        Returns:
            Tupla (ConversationContext, ShortTermMemory)
        """
        self._ensure_initialized()
        
        if not self.isolation_manager:
            return None, ShortTermMemory()
        
        context = self.isolation_manager.get_or_create_context(
            numero_usuario, tipo_conversa, grupo_id
        )
        
        # Carrega short-term memory do contexto
        stm_data = context.short_memory if context else None
        stm = ShortTermMemory(
            conversation_id=context.context_id if context else "",
            context_data={"messages": stm_data} if stm_data else None
        )
        
        return context, stm
    
    def __repr__(self) -> str:
        """Representação textual."""
        return f"ContextBuilder(initialized={self._initialized})"


# ============================================================
# FUNÇÕES DE FÁBRICA
# ============================================================

def criar_context_builder(config_module=None) -> ContextBuilder:
    """
    Factory function para criar ContextBuilder.
    
    Args:
        config_module: Módulo de configuração (opcional)
        
    Returns:
        ContextBuilder instance
    """
    return ContextBuilder(config_module)


# type: ignore