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

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)

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

vectordb = Chroma.from_documents(docs_split, embedding_function)

client = InferenceClient("google/flan-t5-base", token=os.getenv("HUGGINGFACEHUB_API_TOKEN"))

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

def generate_text(context, query):
    inputs = f"Context: {context} Question: {query}"
    response = client.text_generation(prompt=inputs)
    return response
    
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)}")

    # ✅ Corregido: pasar solo lista de strings
    context_strings = [str(c) for c in context]
    reranked = ranker.rank(query, context_strings, 1)
    print(f"🏅 Resultado del reranker: {reranked}")

    # ✅ Seguridad en el acceso al mejor contexto
    best_context = reranked[0].get("text", reranked[0].get("content", context_strings[0]))
    print(f"🧠 Contexto elegido: {best_context[:300]}...")

    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()