"""
07 - Una herramienta que consulta una base de datos de expedientes.

Mismo patrón de siempre; solo cambia el cuerpo de la herramienta: en vez de
internet, consulta una BD. Aquí es SQLite de demo (expedientes.db); en producción
se cambia SOLO la función de conexión por un Oracle/Postgres real (ver NOTA).

Seguridad: NO se deja que el modelo escriba SQL libre. Se exponen consultas
ACOTADAS con parámetros (los filtros usan enum). Así el modelo elige el QUÉ,
pero el CÓMO (el SQL) lo controla el código. Es la regla de oro: prompt sugiere,
código garantiza.

Requisitos:
  - .venv/bin/python crear_expedientes_db.py   (crea la BD de demo)
  - Ollama corriendo con qwen2.5:14b-instruct-q4_K_M
"""

import json
import re
import sqlite3
from pathlib import Path

import ollama

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

ESTADOS = ["abierto", "en_tramite", "cerrado", "archivado"]
TIPOS = ["inspeccion", "licencia", "incidente", "sancion"]


# --- NOTA: para una BD real, cambia SOLO esta función --------------------------
# def _conectar():
#     import oracledb            # o psycopg2 para Postgres
#     return oracledb.connect(user=..., password=..., dsn=...)
# y ajusta los nombres de tabla/columnas en las consultas de abajo.
def _conectar():
    con = sqlite3.connect(DB)
    con.row_factory = sqlite3.Row
    return con


# --- Herramienta 1: buscar con filtros (SQL parametrizado, seguro) -------------

def buscar_expedientes(estado: str = None, tipo: str = None,
                       anio: int = None, responsable: str = None,
                       limite: int = 10) -> str:
    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
    )


# --- Herramienta 2: detalle de un expediente concreto --------------------------

def detalle_expediente(codigo: str) -> str:
    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)


FUNCIONES = {"buscar_expedientes": buscar_expedientes, "detalle_expediente": detalle_expediente}

TOOLS = [
    {"type": "function", "function": {
        "name": "buscar_expedientes",
        "description": "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'.",
        "parameters": {"type": "object", "properties": {
            "estado": {"type": "string", "enum": ESTADOS, "description": "Estado del expediente."},
            "tipo": {"type": "string", "enum": TIPOS, "description": "Tipo de expediente."},
            "anio": {"type": "integer", "description": "Año de apertura, ej. 2025."},
            "responsable": {"type": "string", "description": "Nombre (o parte) del responsable."},
            "limite": {"type": "integer", "description": "Máximo de resultados (por defecto 10)."},
        }}}},
    {"type": "function", "function": {
        "name": "detalle_expediente",
        "description": "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.",
        "parameters": {"type": "object", "properties": {
            "codigo": {"type": "string", "description": "Código del expediente, ej. EXP-2025-008."},
        }, "required": ["codigo"]}}},
]

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


def responder(pregunta: str, max_turnos: int = 6):
    print(f"\n{'='*70}\nPREGUNTA: {pregunta}\n{'='*70}")
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": pregunta},
    ]
    for _ in range(max_turnos):
        r = ollama.chat(model=MODELO, messages=messages, tools=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"   --> {nombre}({dict(args)})")
            f = FUNCIONES.get(nombre)
            res = f(**args) if f else f"Error: '{nombre}' desconocida"
            messages.append({"role": "tool", "name": nombre, "content": str(res)})


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