"""
08 - El agente de expedientes, ahora hablando MCP.

Mismo bucle de agente que 07, pero con DOS cambios que son toda la gracia de MCP:

  (A) El catálogo TOOLS ya NO se escribe a mano. Se DESCUBRE preguntándole al
      servidor (session.list_tools()). El servidor decide qué ofrece.
  (B) La ejecución ya NO usa el dict FUNCIONES local. Se delega al servidor por
      el protocolo (session.call_tool(nombre, args)).

El LLM y el bucle de decisión son idénticos a 07. Lo único que cambió es DE DÓNDE
vienen las herramientas y CÓMO se ejecutan. Eso es, exactamente, lo que aporta MCP.

Es asíncrono porque el cliente MCP lo es. Lanza el servidor como subproceso y
habla con él por stdio.

Requisitos:
  - .venv/bin/python crear_expedientes_db.py   (la BD que usa el servidor)
  - Ollama con qwen2.5:14b-instruct-q4_K_M
"""

import asyncio
import json
import re
import sys
from pathlib import Path

import ollama
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

MODELO = "qwen2.5:14b-instruct-q4_K_M"
SERVER = Path(__file__).parent / "mcp_expedientes_server.py"

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."
)


def llamadas_fugadas_en_texto(contenido: str):
    if not contenido:
        return []
    out = []
    for b in re.findall(r'\{[^{}]*"name"[^{}]*\{[^{}]*\}[^{}]*\}', contenido):
        try:
            o = json.loads(b)
            if "name" in o and "arguments" in o:
                out.append((o["name"], o["arguments"]))
        except json.JSONDecodeError:
            continue
    return out


async def responder(session: ClientSession, ollama_tools, pregunta: str, max_turnos=6):
    print(f"\n{'='*70}\nPREGUNTA: {pregunta}\n{'='*70}")
    cliente = ollama.AsyncClient()
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": pregunta},
    ]
    for _ in range(max_turnos):
        r = await cliente.chat(model=MODELO, messages=messages,
                               tools=ollama_tools, options={"temperature": 0})
        messages.append(r.message)
        pend = [(tc.function.name, tc.function.arguments) for tc in (r.message.tool_calls or [])]
        if not pend:
            pend = llamadas_fugadas_en_texto(r.message.content)
        if not pend:
            print(f"\nRESPUESTA FINAL:\n{r.message.content}")
            return
        for nombre, args in pend:
            print(f"   --> (MCP) call_tool: {nombre}({dict(args)})")
            # (B) Ejecutar la herramienta YA NO es FUNCIONES[nombre](**args):
            #     se delega al servidor MCP por el protocolo.
            resultado = await session.call_tool(nombre, dict(args))
            texto = "\n".join(c.text for c in resultado.content if hasattr(c, "text"))
            messages.append({"role": "tool", "name": nombre, "content": texto})


async def main():
    # Se lanza el servidor MCP como subproceso y se conecta por stdio.
    params = StdioServerParameters(command=sys.executable, args=[str(SERVER)])
    async with stdio_client(params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()

            # (A) DESCUBRIMIENTO: se pregunta al servidor qué herramientas tiene.
            #     Esto sustituye al dict TOOLS escrito a mano de 07.
            lista = await session.list_tools()
            print("Herramientas DESCUBIERTAS en el servidor MCP:")
            for t in lista.tools:
                print(f"  - {t.name}: {t.description.splitlines()[0]}")

            # Se adapta el esquema MCP al formato que espera Ollama.
            # (inputSchema ya es JSON Schema estándar -> encaja directo)
            ollama_tools = [
                {"type": "function", "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.inputSchema,
                }}
                for t in lista.tools
            ]

            await responder(session, ollama_tools, "¿Qué expedientes de tipo incidente están abiertos?")
            await responder(session, ollama_tools, "¿Cuántos expedientes lleva Marta Ruiz y de qué tipos?")
            await responder(session, ollama_tools, "Dame el detalle del expediente EXP-2025-008.")


if __name__ == "__main__":
    asyncio.run(main())