import gradio as gr
from langchain_community.llms import HuggingFaceHub
from langchain_core.output_parsers import StrOutputParser
from langchain import hub
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from huggingface_hub import InferenceClient
from rerankers import Reranker
import os

embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
persist_directory = "db"

if not os.path.exists(persist_directory):
    loader = PyPDFLoader("Constitucion_española.pdf")
    documents = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    docs_split = text_splitter.split_documents(documents)

    vectordb = Chroma.from_documents(docs_split, embedding_function, persist_directory=persist_directory)
    vectordb.persist()
else:
    vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function)

client = InferenceClient("mrm8488/spanish-t5-small-sqac-for-qa", token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))

ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type='colbert')

def generate_text(context, query):
    prompt = f"""Responde en español. Basándote en el siguiente contexto contesta de forma clara y directa.

    Contexto:
    {context}
    
    Pregunta: {query}
    Respuesta:"""
    response = client.text_generation(prompt=prompt, max_new_tokens=200)
    return response.strip()
    
def test_rag_reranking(query, ranker):
    print(f"\n🔍 Pregunta recibida: {query}")

    docs = vectordb.similarity_search_with_score(query)
    print(f"🔎 Documentos recuperados: {len(docs)}")

    context = []
    for i, (doc, score) in enumerate(docs):
        print(f"📄 Doc {i} - Score: {score}")
        if score < 7:
            doc_details = doc.to_json()['kwargs']
            content = doc_details['page_content']
            context.append(content)
            print(f"✅ Doc {i} agregado al contexto")

    if not context:
        print("❌ No se encontró contexto relevante.")
        return "No se encontró información suficiente para responder."

    print(f"📚 Contextos pasados al ranker: {len(context)}")

    context_strings = [str(c) for c in context]
    #print(help(Reranker.models.ColBERTRanker.rank))
    reranked = ranker.rank(query=query, docs=context_strings)
    
    print(f"🏅 Resultado del reranker: {reranked}")

    best_context = reranked[0].document.text  
    print(f"🧠 Contexto elegido: {best_context[:500]}...")

    respuesta = generate_text(best_context, query)
    print(f"💬 Respuesta generada: {respuesta}")

    return respuesta

def responder_chat(message, history):
    respuesta = test_rag_reranking(message, ranker)
    return respuesta

demo = gr.ChatInterface(
    fn=responder_chat,
    title="Chatbot sobre la constitución española",
    theme="soft"
)

if __name__ == "__main__":
    demo.launch()