#!/usr/bin/env python3
"""Comprueba que el RAG recupera de tus notas, con preguntas de prueba.

    flatpak run --command=python3 io.github.qwersyk.Newelle//master \
        ~/COFRE/CODERS/JARVIS/config/probar_rag.py

Lanza unas cuantas consultas contra el indice ya construido y enseña de que
ficheros saca la respuesta. Sirve para ver si el corpus responde bien antes de
fiarse de el, y para decidir si merece ampliarlo con los .txt o dejarlo en solo
notas.
"""
import os
import sys
import time

sys.path.insert(0, "/app/share/newelle")
sys.path.insert(0, os.path.expanduser(
    "~/.var/app/io.github.qwersyk.Newelle/config/pip"))

from gi.repository import Gio  # noqa: E402

from newelle.handlers.rag.llamaindex_handler import LlamaIndexHanlder  # noqa: E402
from newelle.handlers.embeddings.model2vec import Model2VecHandler  # noqa: E402

MODELOS = os.path.expanduser("~/.var/app/io.github.qwersyk.Newelle/config/models")

# Preguntas de dominio: si el corpus sirve, deberian traer notas relevantes.
PREGUNTAS = [
    "como enumerar un host con nmap",
    "escalada de privilegios en linux",
    "que es oasis y como funciona el pub",
    "reverse shell",
]


def main():
    settings = Gio.Settings.new("io.github.qwersyk.Newelle")
    rag = LlamaIndexHanlder(settings, MODELOS)
    if not rag.is_installed():
        print("faltan dependencias del RAG")
        return 1

    rag.embedding = Model2VecHandler(settings, MODELOS)

    class SinModelo:
        def load_model(self, *a, **k):
            return True
    rag.llm = SinModelo()
    rag.secondary_llm = SinModelo()

    print("cargando el indice...", flush=True)
    try:
        rag.load_index()
    except Exception as e:
        print(f"no se pudo cargar el indice: {e}")
        return 1

    for pregunta in PREGUNTAS:
        print(f"\n### {pregunta}")
        ini = time.time()
        try:
            trozos = rag.get_context(pregunta, [])
        except Exception as e:
            print(f"  fallo: {e}")
            continue
        ms = (time.time() - ini) * 1000
        if not trozos:
            print(f"  sin resultados ({ms:.0f} ms)")
            continue
        print(f"  {len(trozos)} trozos en {ms:.0f} ms:")
        for t in trozos[:1]:
            # el trozo suele traer la ruta o el nombre del fichero delante
            resumen = " ".join(str(t).split())[:400]
            print(f"    · {resumen}")
    return 0


if __name__ == "__main__":
    sys.exit(main())