import os
import openai
import requests
import streamlit as st

from utils.util import *
from langchain.memory import ConversationBufferMemory

SAVE_DIR = "uploaded_files"
os.makedirs(SAVE_DIR, exist_ok=True)


def init_session_state():
    if "openai_api_key" not in st.session_state:
        st.session_state.openai_api_key = ""
    if "uploaded_files" not in st.session_state:
        st.session_state.uploaded_files = os.listdir(SAVE_DIR)
    if "huggingface_token" not in st.session_state:
        st.session_state.huggingface_token = ""


init_session_state()

st.set_page_config(page_title="RegBotBeta", page_icon="📜🤖")
st.title("Welcome to RegBotBeta2.0")
st.header("Powered by `LlamaIndex🦙`, `Langchain🦜🔗` and `OpenAI API`")


st.set_page_config(page_title="RegBotBeta", page_icon="📜🤖")
st.title("Welcome to California Drinking Water Regulation Chatbot Demo!")
st.header("Built with `LlamaIndex🦙`and `Langchain🦜🔗`")


uploaded_files = st.file_uploader(
    "Upload Files",
    accept_multiple_files=True,
    type=["pdf", "docx", "txt", "csv"],
    label_visibility="hidden",
)

if uploaded_files:
    for file in uploaded_files:
        if file not in st.session_state.uploaded_files:
            # add the file to session state
            st.session_state.uploaded_files.append(file.name)
            # save the file to the sample_data directory
            with open(os.path.join(SAVE_DIR, file.name), "wb") as f:
                f.write(file.getbuffer())
    st.success("File(s) uploaded successfully!")


def delete_file(filename):
    """Delete file from session state and local filesystem."""
    if filename in st.session_state.uploaded_files and os.path.exists(
        os.path.join(SAVE_DIR, filename)
    ):
        st.session_state.uploaded_files.remove(filename)
        os.remove(os.path.join(SAVE_DIR, filename))
        st.success(f"Deleted {filename}!")
        st.rerun()


if st.session_state.uploaded_files:
    st.write("Uploaded Files:")
    for index, filename in enumerate(st.session_state.uploaded_files):
        col1, col2 = st.columns([4, 1])
        col1.write(filename)
        if col2.button("Delete", key=f"delete_{index}"):
            delete_file(filename)