import streamlit as st
import getpass
import os

os.environ["OPENAI_API_KEY"] = st.secrets['OPENAI_API_KEY']  # agregada en la config de hugginface
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = st.secrets['OPENAI_API_KEY']

from transformers import pipeline
from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains import RetrievalQA
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import HuggingFaceDatasetLoader
from langchain_community.embeddings import HuggingFaceEmbeddings

with st.status("Cargando Aplicación...", expanded=False) as status:
    st.write("Cargando Dataset...")
    #Carga de DATASET
    dataset_name = "Waflon/FAQ"
    page_content_column = "respuestas"
    loader = HuggingFaceDatasetLoader(dataset_name, page_content_column)
    data = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    #Transformado a tipo de dato especifico para esto
    docs = text_splitter.split_documents(data)

    st.write("Cargando modelo de Sentence-Transformer...")
    #Modelo QA sentence similarity
    modelPath = 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' #español
    model_kwargs = {'device':'cpu'} # cuda or cpu
    encode_kwargs = {'normalize_embeddings': False}

    st.write("Embeddings...")
    #Embeddings que transforman a vectores densos multidimensionales las preguntas del SII
    embeddings = HuggingFaceEmbeddings(
        model_name=modelPath,     # Ruta a modelo Pre entrenado
        model_kwargs=model_kwargs, # Opciones de configuracion del modelo
        encode_kwargs=encode_kwargs # Opciones de Encoding
    )

    st.write("Base de datos FAISS...")
    #DB y retriever
    db = FAISS.from_documents(docs, embeddings)  # Create a retriever object from the 'db' with a search configuration where it retrieves up to 4 relevant splits/documents.

    st.write("Retriever...")
    retriever = db.as_retriever(search_kwargs={"k": 3})
    status.update(label="App cargada con exito!", state="complete")

    st.write("Prompt template...")
    prompt_template = """Usa los siguientes fragmentos de contextos para responder una pregunta al final. Por favor sigue las siguientes reglas:
    1. Si la pregunta requiere vinculos, por favor retornar solamente las vinculos de los vinculos sin respuesta
    2. Si no sabes la respuesta, no inventes una respuesta. Solamente di **No pude encontrar la respuesta definitiva, pero tal vez quieras ver los siguientes vinculos** y agregalos a la lista de vinculos.
    3. Si encuentras la respuesta, escribe una respuesta concisa y agrega la lista de vinculos que sean usadas **directamente** para derivar la respuesta. Excluye los vinculos que sean irrelevantes al final de la respuesta
    
    {contexto}
    
    Pregunta: {question}
    Respuesta Util:"""
    
    
    QA_CHAIN_PROMPT = PromptTemplate.from_template(prompt_template) # prompt_template defined above
    llm_chain = LLMChain(llm=ChatOpenAI(), prompt=QA_CHAIN_PROMPT, callbacks=None, verbose=True)
    document_prompt = PromptTemplate(
        input_variables=["page_content", "url"],
        template="Contexto:\n{page_content}\nVinculo: {url}",
    )
    combine_documents_chain = StuffDocumentsChain(
        llm_chain=llm_chain,
        document_variable_name="contexto",
        document_prompt=document_prompt,
        callbacks=None,
    )
    qa = RetrievalQA(
        combine_documents_chain=combine_documents_chain,
        callbacks=None,
        verbose=True,
        retriever=retriever,
        return_source_documents=True,
    )

    
    status.update(label="Finalizado!", state="complete")
    
def main(): 
    pipe = pipeline('sentiment-analysis')
    text = st.text_area('Ingrese su pregunta')
    if text:
        out = pipe(text)
        st.write(out)

if __name__ == "__main__":
    main()