"""
09 - El MISMO agente de expedientes, ahora con LangGraph (el framework).

Cierra la comparativa de las tres formas de construir el bucle:
   - 07: a mano, dict TOOLS + FUNCIONES + while.
   - 08: MCP, las tools en un servidor aparte.
   - 09: LangGraph, el bucle como un GRAFO de estado (este archivo).

Mapa de equivalencias con el 'responder()' hecho a mano:

   el código a mano               LangGraph
   ---------------------          --------------------------------
   escribir el dict TOOLS    -->  @tool (genera el esquema del docstring/hints)
   ollama.chat(tools=...)    -->  ChatOllama(...).bind_tools(tools)
   FUNCIONES[nombre](**args) -->  ToolNode(tools)
   el while de responder()   -->  StateGraph con nodos y aristas
   "¿pidió herramienta?"     -->  tools_condition (arista condicional)
   la lista messages         -->  el estado (MessagesState)
   max_turnos                -->  los ciclos del grafo (con checkpoints, etc.)

Conclusión del curso: LangGraph no inventa nada nuevo; formaliza EXACTAMENTE el
bucle que ya se construyó. A cambio de una dependencia pesada, ahorra el
boilerplate y da gratis persistencia, streaming, human-in-the-loop...

Requisitos:
  - .venv/bin/python crear_expedientes_db.py
  - Ollama con qwen2.5:14b-instruct-q4_K_M
"""

import json
import sqlite3
from pathlib import Path
from typing import Annotated, Literal, Optional

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_ollama import ChatOllama
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode, tools_condition

MODELO = "qwen2.5:14b-instruct-q4_K_M"
DB = Path(__file__).parent / "expedientes.db"


def _conectar():
    con = sqlite3.connect(DB)
    con.row_factory = sqlite3.Row
    return con


# --- Herramientas: misma lógica que 07, pero el esquema lo genera @tool --------

@tool
def buscar_expedientes(
    estado: Optional[Literal["abierto", "en_tramite", "cerrado", "archivado"]] = None,
    tipo: Optional[Literal["inspeccion", "licencia", "incidente", "sancion"]] = None,
    anio: Optional[int] = None,
    responsable: Optional[str] = None,
    limite: int = 10,
) -> str:
    """Busca expedientes en la base de datos aplicando filtros opcionales (estado,
    tipo, año de apertura, responsable). Devuelve una lista resumida. Úsala para
    preguntas como 'expedientes abiertos de 2025' o 'inspecciones de Marta'."""
    sql = ("SELECT codigo, titulo, tipo, estado, fecha_apertura, responsable, "
           "instalacion FROM expedientes WHERE 1=1")
    params = []
    if estado:
        sql += " AND estado = ?"; params.append(estado)
    if tipo:
        sql += " AND tipo = ?"; params.append(tipo)
    if anio:
        sql += " AND fecha_apertura LIKE ?"; params.append(f"{anio}-%")
    if responsable:
        sql += " AND responsable LIKE ?"; params.append(f"%{responsable}%")
    sql += " ORDER BY fecha_apertura DESC LIMIT ?"; params.append(min(int(limite), 50))
    con = _conectar()
    filas = con.execute(sql, params).fetchall()
    con.close()
    if not filas:
        return "No se encontraron expedientes con esos criterios."
    return "\n".join(
        f"{f['codigo']} | {f['estado']} | {f['tipo']} | {f['fecha_apertura']} | "
        f"{f['responsable']} | {f['instalacion']} | {f['titulo']}"
        for f in filas
    )


@tool
def detalle_expediente(codigo: str) -> str:
    """Devuelve TODOS los datos de un expediente concreto a partir de su código
    (ej. 'EXP-2025-008'). Úsala cuando ya conoces el código y quieres el detalle."""
    con = _conectar()
    f = con.execute("SELECT * FROM expedientes WHERE codigo = ?", [codigo]).fetchone()
    con.close()
    if not f:
        return f"No existe el expediente {codigo}."
    return json.dumps(dict(f), ensure_ascii=False, indent=2)


TOOLS = [buscar_expedientes, detalle_expediente]

SYSTEM = (
    "Eres un asistente del registro de expedientes. Responde SOLO con datos "
    "obtenidos de las herramientas; nunca inventes expedientes ni datos. "
    "Usa buscar_expedientes para localizar y detalle_expediente para profundizar "
    "en uno concreto. Responde de forma clara y en español."
)

# El modelo, con las herramientas 'enganchadas' (equivale a ollama.chat(tools=...))
llm = ChatOllama(model=MODELO, temperature=0).bind_tools(TOOLS)


# --- El bucle, ahora como GRAFO ------------------------------------------------

def nodo_agente(state: MessagesState):
    """Nodo que llama al modelo. Equivale a la llamada ollama.chat de responder()."""
    return {"messages": [llm.invoke(state["messages"])]}


# Construcción del grafo: dos nodos (agente y tools) y una arista condicional.
grafo = StateGraph(MessagesState)
grafo.add_node("agente", nodo_agente)
grafo.add_node("tools", ToolNode(TOOLS))          # = FUNCIONES[nombre](**args)
grafo.add_edge(START, "agente")
# tools_condition = "¿el último mensaje pidió herramienta? -> tools, si no -> END"
grafo.add_conditional_edges("agente", tools_condition)
grafo.add_edge("tools", "agente")                 # tras ejecutar, vuelve al modelo
app = grafo.compile()


def responder(pregunta: str):
    print(f"\n{'='*70}\nPREGUNTA: {pregunta}\n{'='*70}")
    estado = app.invoke({"messages": [SystemMessage(SYSTEM), HumanMessage(pregunta)]})
    # Traza: se muestran las herramientas que se ejecutaron por el camino
    for m in estado["messages"]:
        for tc in getattr(m, "tool_calls", None) or []:
            print(f"   --> (LangGraph) {tc['name']}({tc['args']})")
    print(f"\nRESPUESTA FINAL:\n{estado['messages'][-1].content}")


if __name__ == "__main__":
    responder("¿Qué expedientes de tipo incidente están abiertos?")
    responder("¿Cuántos expedientes lleva Marta Ruiz y de qué tipos?")
    responder("Dame el detalle del expediente EXP-2025-008.")

    # --- Alternativa "con pilas incluidas": todo el grafo en UNA línea ---------
    # from langgraph.prebuilt import create_react_agent
    # app = create_react_agent(llm_sin_bind, TOOLS, prompt=SYSTEM)
    # Hace exactamente lo de arriba (nodos + arista condicional) por ti.