import json
import torch
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForCausalLM

# ================================
# CARGA DE MODELOS
# ================================

embedding_model = SentenceTransformer("MongoDB/mdbr-leaf-ir")

tokenizer = AutoTokenizer.from_pretrained("PleIAs/Pleias-RAG-350M")
llm_model = AutoModelForCausalLM.from_pretrained("PleIAs/Pleias-RAG-350M")

# ================================
# CARGA DE DOCUMENTOS (FIX AQUÍ)
# ================================

with open("documents.json", "r", encoding="utf-8") as f:
    documents_dict = json.load(f)

# 🔥 IMPORTANTE: convertir dict → lista de textos
documents = list(documents_dict.values())

# ================================
# PRECOMPUTO DE EMBEDDINGS
# ================================

doc_embeddings = embedding_model.encode(documents, convert_to_tensor=True)

# ================================
# RECUPERAR DOCUMENTOS
# ================================

def recuperar_documentos(consulta, top_k=2, umbral=0.4):
    query_embedding = embedding_model.encode(consulta, convert_to_tensor=True)

    similitudes = torch.nn.functional.cosine_similarity(
        query_embedding.unsqueeze(0),
        doc_embeddings,
        dim=1
    )

    indices_ordenados = torch.argsort(similitudes, descending=True)

    resultados = []
    for idx in indices_ordenados:
        if similitudes[idx].item() >= umbral:
            resultados.append(documents[idx])
        if len(resultados) >= top_k:
            break

    return resultados

# ================================
# GENERAR RESPUESTA
# ================================

def generar_respuesta(consulta, documentos_recuperados):
    contexto = " ".join(documentos_recuperados)

    prompt = f"""Answer the question based only on the context provided
Context: {contexto}
Question: {consulta}
Answer:"""

    inputs = tokenizer(prompt, return_tensors="pt")

    outputs = llm_model.generate(
        **inputs,
        max_new_tokens=150
    )

    respuesta = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # eliminar el prompt del output
    respuesta = respuesta.replace(prompt, "").strip()

    return respuesta

# ================================
# FUNCIÓN FINAL
# ================================

def preguntar(consulta, top_k=2, umbral=0.4):
    docs = recuperar_documentos(consulta, top_k, umbral)
    return generar_respuesta(consulta, docs)