import streamlit as st
from database import KodeksProcessor
from chatbot import Chatbot
import os

def initialize_session_state():
    if 'chatbot' not in st.session_state:
        st.session_state.chatbot = Chatbot()
    if 'messages' not in st.session_state:
        st.session_state.messages = []

def main():
    st.title("Asystent Prawny")
    
    initialize_session_state()
    
    # Inicjalizacja bazy danych (jeśli potrzebna)
    if 'db_initialized' not in st.session_state:
        with st.spinner("Inicjalizacja bazy danych..."):
            processor = KodeksProcessor()
            if not os.path.exists("chroma_db"):
                processor.process_all_files("data/kodeksy")
        st.session_state.db_initialized = True
    
    # Przycisk do czyszczenia historii
    if st.sidebar.button("Wyczyść historię"):
        st.session_state.chatbot.clear_history()
        st.session_state.messages = []
        st.rerun()

    # Wyświetlenie historii czatu
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # Input użytkownika
    if prompt := st.chat_input("Zadaj pytanie dotyczące prawa..."):
        # Dodaj pytanie użytkownika do historii
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        with st.chat_message("user"):
            st.markdown(prompt)

        # Wyszukaj odpowiednie fragmenty w bazie
        processor = KodeksProcessor()
        relevant_chunks = processor.search(prompt)
        
        # Wygeneruj odpowiedź
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            
            context = st.session_state.chatbot.generate_context(
                [{"text": doc} for doc in relevant_chunks['documents'][0]]
            )
            
            for response_chunk in st.session_state.chatbot.get_response(prompt, context):
                full_response += response_chunk
                message_placeholder.markdown(full_response + "▌")
            
            message_placeholder.markdown(full_response)
        
        # Dodaj odpowiedź asystenta do historii
        st.session_state.messages.append({"role": "assistant", "content": full_response})

if __name__ == "__main__":
    main()