import os
import gradio as gr
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings, ChatHuggingFace, HuggingFaceEndpoint
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from huggingface_hub import snapshot_download

# ─── DESCARGAR CHROMADB DEL HUB ──────────────────────────────────────────────
HF_USERNAME = "Yagofue"
REPO_NAME = "rag"
CHROMA_DIR = "chroma_db"

if not os.path.exists(CHROMA_DIR):
    snapshot_download(
        repo_id=f"{HF_USERNAME}/{REPO_NAME}",
        repo_type="dataset",
        local_dir=CHROMA_DIR
    )

# ─── CARGAR EMBEDDINGS Y VECTORDB ────────────────────────────────────────────
model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
embeddings = HuggingFaceEmbeddings(model_name=model_name)

vectordb = Chroma(
    persist_directory=CHROMA_DIR,
    embedding_function=embeddings
)

# ─── CONFIGURAR LLM ──────────────────────────────────────────────────────────

llm_endpoint = HuggingFaceEndpoint(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
    task="conversational",
    max_new_tokens=512,
    huggingfacehub_api_token=os.environ["HF_TOKEN"]
)

llm = ChatHuggingFace(llm=llm_endpoint)


prompt = ChatPromptTemplate.from_messages([
    ("system", "Eres un asistente útil. Usa el siguiente contexto para responder la pregunta en español.\n\nContexto: {context}"),
    ("human", "{question}")
])

rag_chain = prompt | llm | StrOutputParser()

# ─── FUNCIÓN RAG ─────────────────────────────────────────────────────────────
def responder(mensaje, historial):
    docs = vectordb.similarity_search_with_score(mensaje)

    context = []
    for doc, score in docs:
        if score < 7:
            context.append(doc.page_content)

    if not context:
        return "No tengo información para responder a esta pregunta."

    respuesta = rag_chain.invoke({
        "context": "\n\n".join(context),
        "question": mensaje
    })
    return respuesta

# ─── INTERFAZ GRADIO ─────────────────────────────────────────────────────────
with gr.Blocks(title="RAG EOI Logroño B2") as demo:
    gr.Markdown("## 📚 Asistente RAG — Guía B2 EOI Logroño")
    gr.Markdown("Pregunta sobre los exámenes, criterios de evaluación y contenidos del B2.")

    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(placeholder="Escribe tu pregunta aquí...", label="Pregunta")

    def chat(mensaje, historial):
        respuesta = responder(mensaje, historial)
        historial.append((mensaje, respuesta))
        return "", historial

    msg.submit(chat, [msg, chatbot], [msg, chatbot])
    gr.Button("Enviar").click(chat, [msg, chatbot], [msg, chatbot])

demo.launch()