import streamlit as st
import logging
from database import KodeksProcessor
from config import DATABASE_DIR
import os

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def initialize_session_state():
    if 'chatbot' not in st.session_state:
        st.session_state.chatbot = None
    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(DATABASE_DIR):
                logger.info(f"Przetwarzanie plików w katalogu: data/kodeksy")
                processor.process_all_files("data/kodeksy")
            else:
                logger.info(f"Baza danych już istnieje w {DATABASE_DIR}")
            st.session_state.db_initialized = True

    # Przycisk do czyszczenia historii
    if st.sidebar.button("Wyczyść historię"):
        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 = "Oto co znalazłem w bazie danych:\n\n"

            for doc, metadata in zip(relevant_chunks['documents'][0], relevant_chunks['metadatas'][0]):
                full_response += f"**Artykuł {metadata['article']}**\n{doc}\n\n"

            message_placeholder.markdown(full_response)

        # Dodaj odpowiedź asystenta do historii
        st.session_state.messages.append({"role": "assistant", "content": full_response})

    # Sekcja debugowania
    if st.sidebar.checkbox("Pokaż informacje debugowania"):
        st.subheader("Informacje debugowania")
        processor = KodeksProcessor()
        doc_count = processor.collection.count()
        st.write(f"Całkowita liczba dokumentów w bazie danych: {doc_count}")
        if st.button("Przetwórz pliki ponownie"):
            processor.process_all_files("data/kodeksy")
            st.success("Przetwarzanie zakończone")
        if st.button("Pokaż wszystkie dokumenty"):
            processor.list_all_documents()

if __name__ == "__main__":
    main()