from typing import TypedDict, List, Optional, Annotated, Dict, Any
from langchain_core.messages import BaseMessage
import operator

def merge_usage(a: Optional[Dict[str, int]], b: Optional[Dict[str, int]]) -> Dict[str, int]:
    if not a:
        return b or {"prompt": 0, "completion": 0, "total": 0}
    if not b:
        return a
    return {
        "prompt": a.get("prompt", 0) + b.get("prompt", 0),
        "completion": a.get("completion", 0) + b.get("completion", 0),
        "total": a.get("total", 0) + b.get("total", 0)
    }

class AgentState(TypedDict):
    """
    Estado del agente para Agrota.
    """
    messages: Annotated[List[BaseMessage], operator.add]
    
    # Intenciones detectadas (Cola de ejecución secuencial)
    intencion: Optional[str]
    intenciones: List[str]
    intenciones_pendientes: List[str] # ["UB_PRESUPUESTO", "CARTERA"]
    
    # Entidades de Agrota
    user_id: Optional[str] # WhatsApp ID/Phone
    user_name: Optional[str]
    user_position: Optional[str]
    
    # Contexto técnico
    ultimo_resultado: Optional[str] # Resultado agregado (último o consolidado)
    resultados_tecnicos: Annotated[List[str], operator.add] # Se mantiene acumulativo
    dax_query: Optional[str]
    mensaje_skill: Optional[str]
    respuesta_final: Optional[str]
    
    # Consumo
    usage: Annotated[Dict[str, int], merge_usage]
    
    # Información interactiva (botones, etc if needed)
    interactive_data: Optional[Dict[str, Any]]
    
    # Orquestación de Pre-cálculo (Caché Java)
    archetype_key: Optional[str]
    dashboard: Optional[str] # Dashboard actual en ejecución
    is_precalculated: Optional[bool]
    
    # Seguridad y Sesión (12 horas)
    auth_validated_at: Optional[float] # Timestamp linux
    is_authenticating: Optional[bool]
    user_cedula: Optional[str]
    
    error: Optional[str]