import os
import streamlit as st
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain_pinecone import PineconeVectorStore

st.set_page_config(page_title="Opositor", page_icon="📚", layout="wide")
st.title("📚 Pregunta al TREBEP")

MODELOS_LLM = {
    "llama-3.3-70b-versatile": "llama-3.3-70b-versatile",
    "openai/gpt-oss-120b": "openai/gpt-oss-120b",
    "moonshotai/kimi-k2-instruct-0905": "moonshotai/kimi-k2-instruct-0905",
}

with st.sidebar:
    st.header("Configuración")
    modelo_llm = st.selectbox("Modelo de lenguaje", MODELOS_LLM)
    k_docs = st.slider("Número de fragmentos recuperados", 2, 8, 4)

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

@st.cache_resource
def setup_qa(modelo_llm: str, k_docs: int):
    groq_api_key = os.getenv("GROQ_API_KEY")
    pinecone_api_key = os.getenv("PINECONE_API_KEY")
    langsmith_api_key = os.getenv("LANGCHAIN_API_KEY")

    if not groq_api_key:
        raise ValueError("Falta GROQ_API_KEY en los Secrets.")
    if not pinecone_api_key:
        raise ValueError("Falta PINECONE_API_KEY en los Secrets.")

    if langsmith_api_key:
        os.environ["LANGCHAIN_TRACING_V2"] = "true"
        os.environ["LANGCHAIN_API_KEY"] = langsmith_api_key
        os.environ["LANGCHAIN_PROJECT"] = "trebep"

    os.environ["GROQ_API_KEY"] = groq_api_key
    os.environ["PINECONE_API_KEY"] = pinecone_api_key

    embedding = HuggingFaceEmbeddings(
        model_name="intfloat/multilingual-e5-base"
    )

    llm = ChatGroq(
        model_name=modelo_llm,
        temperature=0
    )

    vectorstore = PineconeVectorStore(
        index_name="boe-intfloat-multilingual-e5-base",
        namespace="trebep",
        embedding=embedding,
    )

    retriever = vectorstore.as_retriever(search_kwargs={"k": k_docs})

    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
    )
    return qa

pregunta = st.text_input("Haz una pregunta sobre el TREBEP")

if pregunta:
    try:
        qa = setup_qa(modelo_llm, k_docs)
        respuesta = qa.invoke({"query": pregunta})

        result = respuesta.get("result", "").strip()
        docs = respuesta.get("source_documents", [])

        st.subheader("Respuesta")
        st.success(result)

        st.session_state.chat_history.append({
            "pregunta": pregunta,
            "respuesta": result
        })

        with st.expander("Fuentes recuperadas"):
            if docs:
                for i, doc in enumerate(docs, start=1):
                    st.markdown(f"**Fragmento {i}**")
                    st.write(doc.page_content)
                    if doc.metadata:
                        st.caption(str(doc.metadata))
                    st.divider()
            else:
                st.write("No se recuperaron fragmentos.")

    except Exception as e:
        st.error(f"Error al procesar la consulta: {e}")

with st.sidebar:
    st.header("Historial")
    if st.session_state.chat_history:
        for item in reversed(st.session_state.chat_history):
            st.markdown(f"**Pregunta:** {item['pregunta']}")
            st.write(item["respuesta"])
            st.divider()
    else:
        st.write("Todavía no hay preguntas en esta sesión.")