import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer, util

# --- 1. CARGA DE MODELOS ---
# Usamos un modelo para buscar informacion en los documentos
print("Cargando el modelo de busqueda...")
embed_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")

# Usamos un modelo para que nos ayude a redactar la respuesta
print("Cargando el modelo de lenguaje...")
model_id = "microsoft/Phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
llm_model = AutoModelForCausalLM.from_pretrained(model_id)

# --- 2. BASE DE CONOCIMIENTO ---
# Leemos los documentos que tenemos guardados en el archivo JSON
with open("documents.json", "r", encoding="utf-8") as f:
    docs_dict = json.load(f)
    documents = list(docs_dict.values())

# Preparamos los datos para que el ordenador pueda buscar en ellos rapidamente
print("Preparando los documentos para la busqueda...")
docs_embeddings = embed_model.encode(documents, convert_to_tensor=True)

# --- 3. FUNCIONES DEL SISTEMA ---

def recuperar_documentos(consulta, top_k=2, umbral=0.4):
    """ Busca en nuestra base de datos los textos que mas se parecen a la pregunta """
    # Pasamos la pregunta a un formato que el modelo entienda
    query_embedding = embed_model.encode(consulta, convert_to_tensor=True)
    
    # Comparamos la pregunta con todos los documentos para ver cuales coinciden mejor
    cos_scores = util.cos_sim(query_embedding, docs_embeddings)[0]
    
    # Nos quedamos con los mejores resultados segun lo que hayamos configurado
    top_results = torch.topk(cos_scores, k=min(top_k, len(documents)))
    
    final_docs = []
    for score, idx in zip(top_results[0], top_results[1]):
        # Solo guardamos el documento si se parece lo suficiente a la pregunta
        if score >= umbral:
            final_docs.append(documents[idx])
            
    return final_docs

def generar_respuesta(consulta, documentos_recuperados):
    """ Redacta una respuesta usando la pregunta y la informacion encontrada """
    # Si no hemos encontrado nada, usamos un texto por defecto
    contexto = " ".join(documentos_recuperados) if documentos_recuperados else "No se ha encontrado informacion relevante."
    
    # Preparamos las instrucciones para el modelo
    prompt = f"Context: {contexto}\nQuestion: {consulta}\nAnswer:"
    
    # Convertimos el texto a un formato que el modelo pueda procesar
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # El modelo genera el texto de la respuesta
    with torch.no_grad():
        outputs = llm_model.generate(**inputs, max_new_tokens=100)
    
    # Limpiamos el texto generado para quedarnos solo con lo importante
    resultado_total = tokenizer.decode(outputs[0], skip_special_tokens=True)
    respuesta_limpia = resultado_total.split("Answer:")[-1].strip()
        
    return respuesta_limpia

def preguntar(consulta, top_k=2, umbral=0.4):
    """ Funcion principal que busca la informacion y redacta la respuesta """
    docs = recuperar_documentos(consulta, top_k, umbral)
    respuesta = generar_respuesta(consulta, docs)
    return respuesta, docs