modules/context_manager_v2.py

"""
═══════════════════════════════════════════════════════════════════════
CONTEXT MANAGER V2 — ISOLAMENTO ROBUSTO E ESCALÁVEL
═══════════════════════════════════════════════════════════════════════
Sistema de isolamento de contexto com suporte a:
✅ Conversas isoladas por conversation_id (PV vs Grupo vs Reply Chain)
✅ Fluxo contextual (quem falou com quem)
✅ Separação entre "listen" e "direct message"
✅ Cache inteligente com TTL
✅ Escalabilidade para 1000+ usuários simultâneos
═══════════════════════════════════════════════════════════════════════
"""

import hashlib
import threading
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Tuple, Set
from datetime import datetime, timedelta
from enum import Enum
import json

# ═══════════════════════════════════════════════════════════════════════
# 📋 ENUMS & TIPOS
# ═══════════════════════════════════════════════════════════════════════

class MessageType(Enum):
    """Tipo de mensagem na conversa"""
    DIRECT = "direct"           # Direcionada a KIAMIA (requer resposta)
    CONTEXTUAL = "contextual"   # Contexto do grupo (não direcionada)
    REPLY = "reply"             # É resposta a outra mensagem
    GROUP_INFO = "group_info"   # Info do grupo (join/leave/etc)
    SYSTEM = "system"           # Mensagem de sistema


class ContextType(Enum):
    """Tipo de contexto/conversa"""
    PRIVATE = "pv"              # Conversa privada 1-on-1
    GROUP = "group"             # Conversa em grupo
    REPLY_CHAIN = "reply_chain" # Cadeia de respostas (thread)


@dataclass
class Message:
    """Estrutura de mensagem com metadados completos"""
    id: str
    usuario: str
    numero: str
    texto: str
    tipo: MessageType
    timestamp: float
    conversation_id: str
    context_type: ContextType
    
    # Metadados de fluxo
    quoted_message_id: Optional[str] = None
    quoted_author: Optional[str] = None
    quoted_texto: Optional[str] = None
    is_reply_to_akira: bool = False
    is_akira_message: bool = False
    
    # Contexto amplo (apenas para CONTEXTUAL)
    related_users: List[str] = field(default_factory=list)
    topic_hint: str = ""
    
    # Score de relevância para AKIRA (0.0-1.0)
    relevance_score: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        """Converte para dict para serialização"""
        return {
            'id': self.id,
            'usuario': self.usuario,
            'numero': self.numero,
            'texto': self.texto,
            'tipo': self.tipo.value,
            'timestamp': self.timestamp,
            'conversation_id': self.conversation_id,
            'context_type': self.context_type.value,
            'quoted_message_id': self.quoted_message_id,
            'quoted_author': self.quoted_author,
            'is_reply_to_akira': self.is_reply_to_akira,
            'relevance_score': self.relevance_score,
        }


@dataclass
class ConversationContext:
    """Contexto isolado de uma conversa"""
    conversation_id: str
    context_type: ContextType
    usuario: str
    numero: str
    
    # Histórico separado
    direct_messages: List[Message] = field(default_factory=list)      # Direcionadas
    contextual_messages: List[Message] = field(default_factory=list)  # Contexto
    
    # Metadados
    created_at: float = field(default_factory=time.time)
    last_access: float = field(default_factory=time.time)
    last_modified: float = field(default_factory=time.time)
    
    # Cache
    _cache_direct: Optional[List[Message]] = None
    _cache_contextual: Optional[List[Message]] = None
    _cache_combined: Optional[List[Message]] = None
    _cache_timestamp: float = 0
    CACHE_TTL: int = 300  # 5 minutos
    
    # Lock para thread-safety
    _lock: threading.RLock = field(default_factory=threading.RLock)
    
    def adicionar_message(self, msg: Message) -> None:
        """Adiciona mensagem mantendo isolamento"""
        with self._lock:
            if msg.tipo == MessageType.DIRECT:
                self.direct_messages.append(msg)
            elif msg.tipo in (MessageType.CONTEXTUAL, MessageType.GROUP_INFO):
                self.contextual_messages.append(msg)
            
            self.last_modified = time.time()
            self._invalidate_cache()
            
            # Limita tamanho (escalabilidade)
            self._cleanup_old_messages()
    
    def obter_direct_messages(self, limit: int = 50) -> List[Message]:
        """Obtém APENAS mensagens direcionadas a AKIRA"""
        with self._lock:
            return list(self.direct_messages[-limit:])
    
    def obter_contextual_messages(self, limit: int = 100) -> List[Message]:
        """Obtém APENAS mensagens contextuais (fluxo do grupo)"""
        with self._lock:
            return list(self.contextual_messages[-limit:])
    
    def obter_historico_completo(self, limit: int = 100) -> List[Message]:
        """Obtém histórico completo mantendo ordem temporal"""
        with self._lock:
            all_msgs = self.direct_messages + self.contextual_messages
            all_msgs.sort(key=lambda m: m.timestamp)
            return list(all_msgs[-limit:])
    
    def obter_contexto_grupo(self) -> Dict[str, Any]:
        """Retorna contexto amplo do grupo para AKIRA entender o fluxo"""
        with self._lock:
            # Extrai quem está falando com quem
            reply_chains: Dict[str, List[str]] = {}
            participants: Set[str] = set()
            topics: List[str] = []
            
            for msg in self.contextual_messages[-50:]:
                participants.add(msg.usuario)
                if msg.quoted_author:
                    key = f"{msg.usuario} → {msg.quoted_author}"
                    if key not in reply_chains:
                        reply_chains[key] = []
                    reply_chains[key].append(msg.texto[:50])
                if msg.topic_hint:
                    topics.append(msg.topic_hint)
            
            return {
                'participants': list(participants),
                'recent_reply_chains': reply_chains,
                'topics_discussed': list(set(topics)),
                'contextual_msg_count': len(self.contextual_messages),
                'direct_msg_count': len(self.direct_messages),
            }
    
    def _invalidate_cache(self) -> None:
        """Invalida cache"""
        self._cache_direct = None
        self._cache_contextual = None
        self._cache_combined = None
        self._cache_timestamp = 0
    
    def _cleanup_old_messages(self, max_age_days: int = 7) -> None:
        """Remove mensagens muito antigas para evitar memory leak"""
        cutoff_time = time.time() - (max_age_days * 86400)
        
        self.direct_messages = [m for m in self.direct_messages if m.timestamp > cutoff_time]
        self.contextual_messages = [m for m in self.contextual_messages if m.timestamp > cutoff_time]
        
        # Limita também por quantidade
        max_msgs = 500
        if len(self.direct_messages) > max_msgs:
            self.direct_messages = self.direct_messages[-max_msgs:]
        if len(self.contextual_messages) > max_msgs:
            self.contextual_messages = self.contextual_messages[-max_msgs:]


# ═══════════════════════════════════════════════════════════════════════
# 🏗️ CONTEXT MANAGER ROBUSTO
# ═══════════════════════════════════════════════════════════════════════

class ContextManagerV2:
    """
    Gerenciador de contexto isolado com suporte a escalabilidade.
    
    Características:
    ✅ Isolamento por conversation_id
    ✅ Separação DIRETA vs CONTEXTUAL
    ✅ Fluxo de conversa com reply chains
    ✅ Cache inteligente
    ✅ Thread-safe
    ✅ Cleanup automático
    """
    
    _instance: Optional['ContextManagerV2'] = None
    _lock = threading.Lock()
    
    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        
        # Storage principal: conversation_id -> ConversationContext
        self.contexts: Dict[str, ConversationContext] = {}
        
        # Mapeamento rápido: (numero, tipo_conversa) -> conversation_id
        self.id_cache: Dict[Tuple[str, str], str] = {}
        
        # Locks
        self._contexts_lock = threading.RLock()
        self._cache_lock = threading.RLock()
        
        # Cleanup thread
        self.cleanup_thread = threading.Thread(daemon=True, target=self._cleanup_loop)
        self.cleanup_thread.start()
        
        self._initialized = True
    
    # ═══════════════════════════════════════════════════════════════════
    # 📌 GERAÇÃO DE CONVERSATION_ID
    # ═══════════════════════════════════════════════════════════════════
    
    def gerar_conversation_id(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        reply_to_message_id: Optional[str] = None
    ) -> str:
        """
        Gera conversation_id único e determinístico.
        
        Exemplos:
        - PV com João: f"202391978787009|pv"
        - Grupo AKIRA: f"g120363392399993499|grupo"
        - Reply chain: f"msg_12345|reply"
        """
        # Verifica cache primeiro
        cache_key = (numero, tipo_conversa, grupo_id or "")
        with self._cache_lock:
            if cache_key in self.id_cache:
                return self.id_cache[cache_key]
        
        # Gera novo
        if reply_to_message_id:
            # Reply chain tem seu próprio contexto
            raw = f"{reply_to_message_id}|reply"
            conv_id = hashlib.sha256(raw.encode()).hexdigest()[:32]
        elif tipo_conversa == "grupo" and grupo_id:
            raw = f"g{grupo_id}|grupo"
            conv_id = hashlib.sha256(raw.encode()).hexdigest()[:32]
        else:
            # PV
            raw = f"{numero}|pv"
            conv_id = hashlib.sha256(raw.encode()).hexdigest()[:32]
        
        # Cache
        with self._cache_lock:
            self.id_cache[cache_key] = conv_id
        
        return conv_id
    
    # ═══════════════════════════════════════════════════════════════════
    # 🔍 OBTENÇÃO E CRIAÇÃO DE CONTEXTOS
    # ═══════════════════════════════════════════════════════════════════
    
    def obter_ou_criar_contexto(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        usuario: str = "desconhecido"
    ) -> ConversationContext:
        """Obtém contexto existente ou cria novo"""
        conv_id = self.gerar_conversation_id(numero, tipo_conversa, grupo_id)
        
        with self._contexts_lock:
            if conv_id not in self.contexts:
                context_type = ContextType.GROUP if tipo_conversa == "grupo" else ContextType.PRIVATE
                self.contexts[conv_id] = ConversationContext(
                    conversation_id=conv_id,
                    context_type=context_type,
                    usuario=usuario,
                    numero=numero
                )
            else:
                # Atualiza last_access
                self.contexts[conv_id].last_access = time.time()
            
            return self.contexts[conv_id]
    
    # ═══════════════════════════════════════════════════════════════════
    # ✉️ ADICIONAR MENSAGENS
    # ═══════════════════════════════════════════════════════════════════
    
    def adicionar_message_direta(
        self,
        numero: str,
        usuario: str,
        texto: str,
        tipo_conversa: str = "pv",
        grupo_id: Optional[str] = None,
        quoted_author: Optional[str] = None,
        quoted_texto: Optional[str] = None
    ) -> Message:
        """Adiciona mensagem DIRECIONADA a AKIRA"""
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id, usuario)
        
        msg = Message(
            id=self._gerar_message_id(),
            usuario=usuario,
            numero=numero,
            texto=texto,
            tipo=MessageType.DIRECT,
            timestamp=time.time(),
            conversation_id=contexto.conversation_id,
            context_type=contexto.context_type,
            quoted_author=quoted_author,
            quoted_texto=quoted_texto,
            is_reply_to_akira=False,
            relevance_score=1.0  # Máxima relevância
        )
        
        contexto.adicionar_message(msg)
        return msg
    
    def adicionar_message_contextual(
        self,
        numero: str,
        usuario: str,
        texto: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        quoted_author: Optional[str] = None,
        quoted_texto: Optional[str] = None,
        topic_hint: str = ""
    ) -> Message:
        """
        Adiciona mensagem CONTEXTUAL (não direcionada a AKIRA).
        AKIRA vê isso para entender o fluxo: "fulano respondeu a beltrano sobre X"
        """
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id, usuario)
        
        msg = Message(
            id=self._gerar_message_id(),
            usuario=usuario,
            numero=numero,
            texto=texto,
            tipo=MessageType.CONTEXTUAL,
            timestamp=time.time(),
            conversation_id=contexto.conversation_id,
            context_type=contexto.context_type,
            quoted_author=quoted_author,
            quoted_texto=quoted_texto,
            topic_hint=topic_hint,
            relevance_score=0.3  # Relevância baixa (apenas contexto)
        )
        
        contexto.adicionar_message(msg)
        return msg
    
    # ═══════════════════════════════════════════════════════════════════
    # 📖 OBTER HISTÓRICOS
    # ═══════════════════════════════════════════════════════════════════
    
    def obter_historico_direto(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        limit: int = 50
    ) -> List[Message]:
        """Obtém APENAS mensagens direcionadas a AKIRA"""
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id)
        return contexto.obter_direct_messages(limit)
    
    def obter_historico_contextual(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        limit: int = 100
    ) -> List[Message]:
        """Obtém fluxo de conversa (quem falou com quem)"""
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id)
        return contexto.obter_contextual_messages(limit)
    
    def obter_historico_completo(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None,
        limit: int = 100
    ) -> List[Message]:
        """Obtém histórico completo em ordem temporal"""
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id)
        return contexto.obter_historico_completo(limit)
    
    def obter_contexto_grupo_amplificado(
        self,
        numero: str,
        tipo_conversa: str,
        grupo_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """Retorna contexto amplo do grupo para AKIRA"""
        contexto = self.obter_ou_criar_contexto(numero, tipo_conversa, grupo_id)
        return contexto.obter_contexto_grupo()
    
    # ═══════════════════════════════════════════════════════════════════
    # 🧹 LIMPEZA & MANUTENÇÃO
    # ═══════════════════════════════════════════════════════════════════
    
    def _cleanup_loop(self, interval: int = 3600) -> None:
        """Background thread que limpa contextos antigos"""
        while True:
            try:
                time.sleep(interval)
                self._cleanup_old_contexts()
            except Exception as e:
                print(f"❌ Erro no cleanup: {e}")
    
    def _cleanup_old_contexts(self, max_age_days: int = 30) -> None:
        """Remove contextos não acessados há muito tempo"""
        cutoff_time = time.time() - (max_age_days * 86400)
        
        with self._contexts_lock:
            to_remove = [
                conv_id for conv_id, ctx in self.contexts.items()
                if ctx.last_access < cutoff_time
            ]
            
            for conv_id in to_remove:
                del self.contexts[conv_id]
    
    def _gerar_message_id(self) -> str:
        """Gera ID único para mensagem"""
        import uuid
        return str(uuid.uuid4())[:12]
    
    # ═══════════════════════════════════════════════════════════════════
    # 📊 STATS & MONITORAMENTO
    # ═══════════════════════════════════════════════════════════════════
    
    def obter_stats(self) -> Dict[str, Any]:
        """Retorna estatísticas do context manager"""
        with self._contexts_lock:
            total_msgs = sum(
                len(ctx.direct_messages) + len(ctx.contextual_messages)
                for ctx in self.contexts.values()
            )
            
            return {
                'total_contexts': len(self.contexts),
                'total_messages': total_msgs,
                'average_msgs_per_context': total_msgs / max(1, len(self.contexts)),
                'cache_size': len(self.id_cache),
                'memory_estimate_mb': (total_msgs * 0.5) / 1024  # Estimativa bruta
            }


# ═══════════════════════════════════════════════════════════════════════
# 🎯 EXPORTS
# ═══════════════════════════════════════════════════════════════════════

def get_context_manager() -> ContextManagerV2:
    """Obtém instância singleton do context manager"""
    return ContextManagerV2()


__all__ = [
    'ContextManagerV2',
    'ConversationContext',
    'Message',
    'MessageType',
    'ContextType',
    'get_context_manager',
]