import os
import streamlit as st
from bot import PDFChatBot

# Folder untuk menyimpan file yang diunggah
UPLOAD_FOLDER = "uploads/"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Inisialisasi chatbot
pdf_chatbot = PDFChatBot()

# Streamlit Page Configuration
st.set_page_config(
    page_title="RAG Chatbot Q&A",
    layout="wide",
    page_icon="📘",
)

# Sidebar untuk unggah file
with st.sidebar:
    st.header("PDF Upload")
    uploaded_file = st.file_uploader("Upload your PDF", type=["pdf"])
    st.info("Supported file type: PDF")

# Chatbot Title
st.title("RAG Chatbot Q&A")

# Chat Display State
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Simpan file yang diunggah dan proses
if uploaded_file:
    # Simpan file di folder uploads
    file_path = os.path.join(UPLOAD_FOLDER, uploaded_file.name)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Render halaman pertama PDF
    st.header("PDF Overview")
    pdf_image = pdf_chatbot.render_page(file_path)
    st.image(pdf_image, use_column_width=True)

# Chat Interface
st.subheader("Chat with Your PDF")
user_input = st.text_input("Ask a question about your PDF:")
if st.button("Send") and user_input:
    # Tambahkan input user ke chat history
    st.session_state.chat_history.append({"role": "user", "content": user_input})

    if uploaded_file:
        # Proses input user dan buat respon
        with st.spinner("Processing..."):
            response = pdf_chatbot.generate_response(
                query=user_input,
                file=file_path,
            )
        st.session_state.chat_history.append({"role": "assistant", "content": response})

    else:
        st.warning("Please upload a PDF before asking a question!")

# Tampilkan chat history
for chat in st.session_state.chat_history:
    if chat["role"] == "user":
        st.chat_message("user").markdown(chat["content"])
    elif chat["role"] == "assistant":
        st.chat_message("assistant").markdown(chat["content"])