import streamlit as st
import os
import shutil

from agents.customer_service import get_customer_agents_response, init_support_agents
from agents.customer_service_agent import get_customer_service_agent

st.title("📁 Customer Assistant Chatbot")

KNOWLEDGE_FOLDER = "tmp/knowledge-base"
MEMORY_FOLDER = "tmp/memory"
os.makedirs(KNOWLEDGE_FOLDER, exist_ok=True)


def clear_knowledge():
    shutil.rmtree(KNOWLEDGE_FOLDER, ignore_errors=True)
    shutil.rmtree(MEMORY_FOLDER, ignore_errors=True)
    os.makedirs(KNOWLEDGE_FOLDER, exist_ok=True)
    os.makedirs(MEMORY_FOLDER, exist_ok=True)
    # get_customer_service_agent(True, True)
    init_support_agents(True, True)
    st.success("Knowledge base cleared.")


# st.sidebar.button("🗑️ Clear Knowledge Base", on_click=clear_knowledge)

if "file_uploaded" not in st.session_state:
    st.session_state.file_uploaded = None

uploaded_file = st.file_uploader(" ", type=["pdf"], key="uploader")
if uploaded_file and not uploaded_file.file_id == st.session_state.file_uploaded:
    with st.spinner("Sto esaminando il documento..."):
        clear_knowledge()
        file_path = os.path.join(KNOWLEDGE_FOLDER, uploaded_file.name)
        with open(file_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.success(f"Uploaded and stored **{uploaded_file.name}**")
        # get_customer_service_agent(True)
        init_support_agents(True)
        st.session_state.file_uploaded = uploaded_file.file_id
        st.rerun()

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

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Come posso aiutarti?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        with st.spinner("Sto pensando..."):
            # chunks = get_customer_service_agent().run(prompt)
            response, session_id = get_customer_agents_response(
                prompt,
                user_id="user_1",
                session_id=st.session_state.get("session_id"),
                history=st.session_state.messages
            )
        # response = chunks.content
        st.session_state.messages.append({"role": "assistant", "content": response})
        st.rerun()