import streamlit as st
import os

from groq import Groq
from datetime import datetime
from langchain_community.embeddings import HuggingFaceEmbeddings 
from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA
from langchain_pinecone import PineconeVectorStore

st.set_page_config('Opositor')
st.header("Pregunta al trebep")

modelos_llm = [
    'llama3-70b-8192',
    'llama3-8b-8192',
    'mixtral-8x7b-32768',
    'gemma-7b-it'
]
modelo_llm = st.selectbox('Modelo de lenguaje', list(modelos_llm))

@st.cache_resource 
def setup(modelo_llm):
    # Langsmith
    os.environ["LANGCHAIN_TRACING_V2"] = "true"
    os.environ["LANGCHAIN_API_KEY"] = "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"
    os.environ["LANGCHAIN_PROJECT"] = "trebep"

    # CARGAMOS MODELO DE EMBEDDING
    model_name = 'intfloat/multilingual-e5-base'
    embedding = HuggingFaceEmbeddings(model_name=model_name)

    # CARGAMOS LLM
    os.environ["GROQ_API_KEY"] = "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"
    llm = ChatGroq(model = modelo_llm)
    
    # CARGAMOS EL VECTORSTORE DE PINECONE
    os.environ["PINECONE_API_KEY"] ='4bf0b4cf-4ced-4f70-8977-d60bb8ae405a'
    index_name = "boe-intfloat-multilingual-e5-base"
    namespace = "trebep"
    vectorstore = PineconeVectorStore(index_name=index_name, namespace=namespace, embedding=embedding)

    # CREAMOS EL RETRIEVAL
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever(),
        return_source_documents=True,
        #verbose=True
    )

    return qa

# Función para mostrar logs
def mostrar_logs(logs,hints):
    # Crear un contenedor desplegable
    with st.expander("Chunks"):
        for hint in hints:
            st.write(hint.page_content)
            st.write("-" * 30)

    st.sidebar.header("Registro de preguntas")
    for entry in logs:
        st.sidebar.write(f"**Pregunta: {entry['Pregunta']}**")
        st.sidebar.write(f"Respuesta: {entry['Respuesta']}")


# Lista para almacenar preguntas y respuestas
logs = []

user_question = st.text_input("¡A jugar! Haz una pregunta al trebep:")
if user_question:

    # Inicializar entorno
    qa = setup(modelo_llm)

    # Obtenemos la respuesta
    respuesta = qa.invoke(user_question)
    docs = respuesta['source_documents']
    #result = respuesta['result']
    result = respuesta['result'].strip("[]")


    # Mostrar la variable en color verde
    st.subheader("Respuesta")
    st.write(f":green[{str(result)}]")

    # Guardar pregunta y respuesta en los logs
    logs.append({"Pregunta": user_question, "Respuesta": result})

    # Mostrar logs actualizados
    mostrar_logs(logs,docs)