import streamlit as st
from chroma.chatbot_retriever import ask
from chroma.chatbot_chroma_db import create_new_chat, delete_chat, load_chats_from_db

# Initialize session state
if "chats" not in st.session_state:
    loaded = load_chats_from_db()
    if loaded:
        st.session_state.chats = loaded
    else:
        initial_chat_id = create_new_chat()
        st.session_state.chats = {"Chat 1": {"messages": [], "chat_id": initial_chat_id}}

if "active_chat" not in st.session_state:
    st.session_state.active_chat = "Chat 1"


# st.write(list(st.session_state.chats.keys()))

# print(st.session_state.active_chat)


# sidebar
with st.sidebar:
    st.title("Bank Chatbot")
    # New Chat button
    if st.button("+ New Chat", use_container_width=True):
        new_name = f"Chat {len(st.session_state.chats) + 1}"
        new_chat_id = create_new_chat()
        st.session_state.chats[new_name] = {"messages": [], "chat_id": new_chat_id}
        st.session_state.active_chat = new_name
        st.rerun()

    st.divider()

    # List all chats as buttons
    for chat_name in list(st.session_state.chats.keys()):  # list() avoids mutation during loop
        col1, col2 = st.columns([4, 1])

        with col1:
            if st.button(chat_name, key=f"select_{chat_name}", use_container_width=True):
                st.session_state.active_chat = chat_name
                st.rerun()

        with col2:
            if st.button("🗑", key=f"delete_{chat_name}"):
                delete_chat(st.session_state.chats[chat_name]["chat_id"])

                del st.session_state.chats[chat_name]

                # if deleted chat was active, switch to another
                if st.session_state.active_chat == chat_name:
                    remaining = list(st.session_state.chats.keys())
                    if remaining:
                        st.session_state.active_chat = remaining[0]
                    else:
                        # no chats left, create a fresh one
                        new_chat_id = create_new_chat()

                        st.session_state.chats["Chat 1"] = {
                            "messages": [],
                            "chat_id": new_chat_id
                        }
                        st.session_state.active_chat = "Chat 1"

                st.rerun()

st.header(st.session_state.active_chat)

active = st.session_state.active_chat
# st.write(st.session_state.chats)
# Display history for the active chat
for msg in st.session_state.chats[active]["messages"]:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])

if prompt := st.chat_input("Ask me anything..."):
    # if "chat_id" not in st.session_state:
    #     st.session_state.chat_id = create_new_chat()

    active = st.session_state.active_chat
    chat_id = st.session_state.chats[active]["chat_id"]
    
    # saves user message
    st.session_state.chats[active]["messages"].append({"role": "user", "content": prompt})
    
    # with st.chat_message("user"):
    #     st.write(prompt)

    # add response from RAG chain, passing session chat history
    # with st.chat_message("assistant"):
    with st.spinner("Thinking..."):
        response = ask(prompt, chat_id = chat_id, chat_name=active)
    # st.write(response['answer'])

    
    st.session_state.chats[active]["messages"].append({"role": "assistant", "content": response})
    st.rerun()