import streamlit as st
import os
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import HuggingFaceHub
from langchain.chains.question_answering import load_qa_chain
from langchain.prompts import PromptTemplate

st.set_page_config(page_title='preguntaDOC')
st.header("Pregunta a tu PDF")

# Ya no necesitamos la clave de OpenAI
huggingface_api_token = st.text_input('Hugging Face API Token (opcional)', type='password')

pdf_obj = st.file_uploader("Carga tu documento", type="pdf", on_change=st.cache_resource.clear)

@st.cache_resource 
def create_embeddings(pdf):
    pdf_reader = PdfReader(pdf)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100,
        length_function=len
    )        
    chunks = text_splitter.split_text(text)
    
    # Usando el mismo modelo de embeddings que ya estabas usando
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
    knowledge_base = FAISS.from_texts(chunks, embeddings)
    return knowledge_base

if pdf_obj:
    knowledge_base = create_embeddings(pdf_obj)
    user_question = st.text_input("Haz una pregunta sobre tu PDF:")
    
    if user_question:
        # Configurar el token de Hugging Face si se proporciona
        if huggingface_api_token:
            os.environ["HUGGINGFACEHUB_API_TOKEN"] = huggingface_api_token
        
        docs = knowledge_base.similarity_search(user_question, 3)
        
        # Usar un modelo gratuito de Hugging Face en lugar de OpenAI
        llm = HuggingFaceHub(
            repo_id="google/flan-t5-large",  # Modelo gratuito con buenas capacidades para Q&A
            model_kwargs={"temperature": 0.5, "max_length": 512}
        )
        
        # Crear un prompt template adecuado para modelos como T5
        prompt_template = """
        Responde a la siguiente pregunta basándote únicamente en el contexto proporcionado.
        
        Contexto: {context}
        
        Pregunta: {question}
        
        Respuesta:
        """
        
        PROMPT = PromptTemplate(
            template=prompt_template, 
            input_variables=["context", "question"]
        )
        
        chain = load_qa_chain(llm, chain_type="stuff", prompt=PROMPT)
        
        # Mostrar un mensaje mientras se procesa
        with st.spinner("Procesando tu pregunta..."):
            try:
                respuesta = chain.run(input_documents=docs, question=user_question)
                st.write(respuesta)
            except Exception as e:
                st.error(f"Error al procesar tu pregunta: {str(e)}")
                st.info("Si no has proporcionado un token de Hugging Face, considera hacerlo para evitar limitaciones de rate limit.")