import streamlit as st
from openai import OpenAI
from langchain.document_loaders import UnstructuredFileLoader

st.title("Extração de resultados")


def load_document(file_path, mode="single"):
    loader = UnstructuredFileLoader(file_path, mode=mode)
    docs = loader.load()

    if mode == "single":
        return docs[0].page_content, []
    elif mode == "elements":
        return " ".join([str(doc.page_content) for doc in docs]), docs


# Sidebar para entrada de chave API e seleção de modelo
with st.sidebar:
    openai_api_key = st.text_input(
        "OpenAI API Key", key="chatbot_api_key", type="password"
    )
    st.markdown(
        "[Pegue aqui sua chave OpenAI API](https://platform.openai.com/account/api-keys)"
    )

    uploaded_file = st.file_uploader(
        "Escolha um arquivo PDF ou imagem", type=["pdf", "png", "jpg", "jpeg", "docx"]
    )

    if st.button("Limpar Conversa"):
        st.session_state["messages"] = [
            {"role": "assistant", "content": "How can I help you?"}
        ]
        st.experimental_rerun()

st.title("💬 Chatbot")
st.caption("🚀 A Streamlit chatbot powered by OpenAI")

models = "asst_wXmB8syPqCxPWGtNqiuOXQtn"

if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "assistant", "content": "How can I help you?"}
    ]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])


if prompt := st.chat_input():
    if not openai_api_key:
        st.info("Insira sua chave API OpenAI para continuar.")
        st.stop()

    client = OpenAI(api_key=openai_api_key)
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    response = client.chat.completions.create(
        model=models, messages=st.session_state.messages
    )
    msg = response.choices[0].message.content
    st.session_state.messages.append({"role": "assistant", "content": msg})
    st.chat_message("assistant").write(msg)


if uploaded_file is not None:
    data = load_document(st.file_uploader)
    if data := st.session_state.get("uploaded_file", None) is not None and st.button(
        "Enviar", type="primary"
    ):
        if data is not None:
            st.session_state.exames({"role": "uploaded_file", "content": data})
            resposta = client.chat.completions.create(
                model=models, messages=st.session_state.exames
            )
            msg = resposta.choices[0].message.content
            st.session_state.exames.append({"role": "assistant", "content": msg})
            st.chat_message("assistant").write(msg)
        else:
            st.error("Não foi possível extrair texto do arquivo carregado.")
    else:
        st.warning("Por favor, carregue um arquivo antes de enviar.")