Agente de Soporte Multicanal
con LangGraph + Estado Persistente

# ============================================================
# LearnSpark — Agente de Soporte con LangGraph
# Cliente: LearnSpark SaaS LMS | CULTIVA IA
# Patron: Supervisor + 4 especialistas + Human-in-the-Loop
# Persistencia: SqliteSaver (checkpoints.db)
# ============================================================

from typing import Annotated, TypedDict, Optional, Literal
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.sqlite import SqliteSaver
from operator import add

# ── Estado compartido ─────────────────────────────────────────
class SupportState(TypedDict):
    messages:         Annotated[list, add_messages]
    customer_id:      str
    ticket_category:  str              # billing | users | integrations | courses
    urgency:          str              # low | medium | high
    tool_calls_made:  Annotated[list[str], add]
    resolved:         bool
    iteration_count:  int
    escalation_reason:Optional[str]
    next_agent:       str

# ── Herramientas de cada especialista ─────────────────────────
@tool
def lookup_subscription(customer_id: str) -> str:
    """Consulta el estado de suscripcion, plan y fecha de renovacion."""
    # En produccion: llamada a Stripe API o base de datos LearnSpark
    return f"Plan Scale | 500 seats | Renovacion: 2025-01-15 | Estado: ACTIVO"

@tool
def get_invoice_pdf(invoice_id: str) -> str:
    """Genera URL de descarga para factura especifica."""
    return f"https://billing.learnspark.io/invoices/{invoice_id}.pdf?token=abc123"

@tool
def invite_user(email: str, role: str = "member") -> str:
    """Envia invitacion a nuevo usuario con rol especificado."""
    return f"Invitacion enviada a {email} con rol '{role}'. Expira en 72h."

@tool
def list_team_members(customer_id: str) -> str:
    """Lista miembros del equipo con estado y ultimo acceso."""
    return "432/500 seats usados | 12 pendientes de activar | 8 inactivos (+90d)"

@tool
def check_oauth_status(integration: str) -> str:
    """Verifica estado de integracion OAuth (Slack, Jira, Google Workspace)."""
    return f"Integracion {integration}: token expirado el 2024-11-28. Requiere re-autenticacion."

@tool
def reset_api_token(customer_id: str, service: str) -> str:
    """Regenera el API token para una integracion especifica."""
    return f"Token regenerado para {service}. Nuevo token: ls_live_xxx...yyy (visible 1 sola vez)"

@tool
def check_video_cdn(course_id: str) -> str:
    """Comprueba disponibilidad del video en CDN y calidad de stream."""
    return f"Curso {course_id}: CDN OK | Latencia: 24ms | Codec: H.264 | Bitrate adaptativo: activo"

# ── LLM base ──────────────────────────────────────────────────
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# ── Nodo 1: Supervisor (clasificador y router) ────────────────
def supervisor_node(state: SupportState) -> dict:
    response = llm.invoke([
        ("system", """Eres el supervisor de soporte de LearnSpark.
Clasifica el ticket y responde SOLO con JSON:
{
  "category": "billing|users|integrations|courses|unknown",
  "urgency": "low|medium|high",
  "next_agent": "billing|users|integrations|courses|human",
  "reasoning": "breve explicacion"
}
Escala a 'human' si: urgency=high + no se puede resolver con herramientas."""),
        *state["messages"]
    ])
    # Parse JSON response (simplificado para el ejemplo)
    import json, re
    match = re.search(r'\{.*\}', response.content, re.DOTALL)
    data = json.loads(match.group()) if match else {}
    return {
        "messages": [response],
        "ticket_category": data.get("category", "unknown"),
        "urgency": data.get("urgency", "medium"),
        "next_agent": data.get("next_agent", "human"),
    }

# ── Nodo 2: Agentes especialistas (patron ReAct con tools) ────
def make_specialist(name: str, system_prompt: str, tools: list):
    specialist_llm = llm.bind_tools(tools)
    def node(state: SupportState) -> dict:
        if state["iteration_count"] >= 5:
            return {"next_agent": "human", "escalation_reason": "Max iteraciones alcanzado"}
        response = specialist_llm.invoke([("system", system_prompt), *state["messages"]])
        resolved = not response.tool_calls and "resuelto" in response.content.lower()
        return {
            "messages": [response],
            "iteration_count": state["iteration_count"] + 1,
            "resolved": resolved,
            "tool_calls_made": [tc["name"] for tc in (response.tool_calls or [])],
        }
    node.__name__ = name
    return node

# ── Nodo 3: Human escalation (HITL) ──────────────────────────
def human_escalation_node(state: SupportState) -> dict:
    # interrupt_before pausa aqui — el agente espera aprobacion
    return {"messages": [("assistant",
        f"Ticket escalado a soporte humano. Motivo: {state.get('escalation_reason', 'Complejidad alta')}. "
        f"SLA: respuesta en < 2h para urgency={state['urgency']}. Ticket #LS-{state['customer_id'][-4:]}-H"
    )]}

# ── Ensamblaje del grafo ──────────────────────────────────────
billing_tools    = [lookup_subscription, get_invoice_pdf]
users_tools      = [invite_user, list_team_members]
integration_tools= [check_oauth_status, reset_api_token]
courses_tools    = [check_video_cdn]

graph = StateGraph(SupportState)

# Nodos
graph.add_node("supervisor",    supervisor_node)
graph.add_node("billing",       make_specialist("billing", "Especialista en facturacion LearnSpark.", billing_tools))
graph.add_node("users",         make_specialist("users", "Especialista en gestion de equipos.", users_tools))
graph.add_node("integrations",  make_specialist("integrations", "Especialista en OAuth e integraciones.", integration_tools))
graph.add_node("courses",       make_specialist("courses", "Especialista en reproduccion de cursos.", courses_tools))
graph.add_node("tools",         ToolNode([*billing_tools, *users_tools, *integration_tools, *courses_tools]))
graph.add_node("human",         human_escalation_node)

# Aristas
graph.add_edge(START, "supervisor")
graph.add_conditional_edges("supervisor",
    lambda s: s["next_agent"],
    {"billing": "billing", "users": "users",
     "integrations": "integrations", "courses": "courses", "human": "human"}
)
for specialist in ["billing", "users", "integrations", "courses"]:
    graph.add_conditional_edges(specialist,
        lambda s: "tools" if s["messages"][-1].tool_calls else ("human" if s.get("next_agent") == "human" else END),
        {"tools": "tools", "human": "human", END: END}
    )
    graph.add_edge("tools", specialist)
graph.add_edge("human", END)

# Compilar con persistencia SQLite + HITL antes de escalado humano
memory = SqliteSaver.from_conn_string("learnspark_checkpoints.db")
app    = graph.compile(checkpointer=memory, interrupt_before=["human"])

# ── Uso: streaming por nodos ──────────────────────────────────
config = {"configurable": {"thread_id": "learnspark_tenant_4821"}}

ticket = {
    "messages": [("human", "No puedo conectar Jira. El token parece expirado y estamos bloqueados.")],
    "customer_id": "learnspark_tenant_4821",
    "ticket_category": "",
    "urgency": "",
    "tool_calls_made": [],
    "resolved": False,
    "iteration_count": 0,
    "escalation_reason": None,
    "next_agent": "",
}

for chunk in app.stream(ticket, config, stream_mode="updates"):
    for node_name, output in chunk.items():
        print(f"[{node_name}]", output.get("messages", [{}])[-1])