app.py

import os
import logging
import streamlit as st
from typing import Dict, Any

from config import Config
from backend.document_service import DocumentService
from frontend.styles import CUSTOM_CSS
from frontend.sidebar import render_sidebar
from frontend.chat_view import render_chat_view
from frontend.library_view import render_library_view
from frontend.audit_view import render_audit_view
from utils.logger import setup_logger

# Initialize Logger
logger = setup_logger(__name__)

# Page Config
st.set_page_config(
    page_title=Config.APP_TITLE,
    page_icon=Config.APP_ICON,
    layout=Config.LAYOUT,
    initial_sidebar_state="expanded",
)

# Inject CSS
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)

@st.dialog("Import Document")
def render_upload_dialog(document_service: DocumentService) -> None:
    with st.form("upload_form", clear_on_submit=True):
        uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
        doc_type = st.selectbox("Document Type", ("General Document", "Rule Set"))
        
        submitted = st.form_submit_button("Add to Library")
        
        if submitted and uploaded_file:
            _handle_upload(uploaded_file, doc_type, document_service)

def _handle_upload(uploaded_file: Any, doc_type: str, document_service: DocumentService) -> None:
    try:
        temp_path = os.path.join(Config.TEMP_DIR, uploaded_file.name)
        # Ensure temp dir exists
        if not os.path.exists(Config.TEMP_DIR):
            os.makedirs(Config.TEMP_DIR)

        with open(temp_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        
        with st.spinner(f"Processing {uploaded_file.name}..."):
            doc, status = document_service.process_document(temp_path, doc_type)
            
            if status == 'unchanged':
                st.toast(f"'{doc['filename']}' is already up to date.", icon="✅")
                logger.info(f"Document unchanged: {doc['filename']}")
            else:
                st.toast(f"'{doc['filename']}' processed successfully!", icon="🎉")
                logger.info(f"Document processed: {doc['filename']}")
            
            st.rerun()
    except Exception as e:
        error_msg = f"Error processing document: {e}"
        st.error(error_msg)
        logger.error(error_msg)

@st.dialog("Delete Document")
def render_delete_dialog(doc: Dict[str, Any], document_service: DocumentService) -> None:
    st.warning(f"Are you sure you want to delete '{doc.get('filename', 'this document')}'?")
    
    if st.button("Yes, Delete", type="primary"):
        try:
            if document_service.delete_document(doc['doc_id']):
                st.toast("Document deleted.")
                logger.info(f"Document deleted: {doc['doc_id']}")
                st.rerun()
            else:
                st.error("Failed to delete document.")
                logger.error(f"Failed to delete document: {doc['doc_id']}")
        except Exception as e:
            st.error(f"Error during deletion: {e}")
            logger.error(f"Error deleting document {doc['doc_id']}: {e}")

import time

def render_settings_view() -> None:
    st.header("Settings")
    st.markdown("---")
    
    st.subheader("AI Configuration")
    llm_provider = st.selectbox(
        "LLM Provider",
        ("Ollama (Private Local Inference)", "Gemini (Cloud Performance)"),
        key='llm_provider_select'
    )
    
    if st.session_state.get('llm_provider') != llm_provider:
        st.session_state.llm_provider = llm_provider
        logger.info(f"LLM Provider switched to: {llm_provider}")
        st.rerun()
        
    st.info(f"Current Provider: {llm_provider}")

    st.markdown("---")
    st.subheader("Maintenance")
    st.caption("Removes data from the knowledge base that is no longer linked to any document in your library.")
    
    if st.button("Clean Up Knowledge Base", type="primary"):
        with st.spinner("Scanning for orphaned vectors..."):
            try:
                # Get valid IDs
                docs = st.session_state.document_service.get_all_documents()
                valid_ids = [d['doc_id'] for d in docs]
                
                # Perform purge
                count = st.session_state.document_service.vector_store_service.purge_orphaned_vectors(valid_ids)
                
                if count > 0:
                    st.success(f"Cleaned up {count} orphaned document fragments.")
                    time.sleep(2) # Let user see the success message
                    st.rerun()
                else:
                    st.info("Database is healthy. No orphans found.")
            except Exception as e:
                st.error(f"Repair failed: {e}")
                logger.error(f"Database repair failed: {e}")

def main() -> None:
    # Initialize Session State
    if 'llm_provider' not in st.session_state:
        st.session_state.llm_provider = "Ollama (Private Local Inference)"

    # Initialize Services
    if 'document_service' not in st.session_state:
        Config.ensure_dirs()
        with st.spinner("Initializing AI Models & Database..."):
            try:
                st.session_state.document_service = DocumentService()
                logger.info("DocumentService initialized successfully.")
            except Exception as e:
                st.error(f"Critical Error: Failed to initialize services. {e}")
                logger.critical(f"Service initialization failed: {e}")
                st.stop()

    # Render Sidebar
    render_sidebar(
        st.session_state.document_service,
        on_upload_click=lambda: render_upload_dialog(st.session_state.document_service),
        on_delete_click=lambda doc: render_delete_dialog(doc, st.session_state.document_service)
    )

    # Main Layout Routing
    current_view = st.session_state.get('current_view', 'chat')
    
    if current_view == 'chat':
        render_chat_view()
    elif current_view == 'library':
        render_library_view(st.session_state.document_service, on_audit_click=None)
    elif current_view == 'audit':
        render_audit_view(st.session_state.document_service)
    elif current_view == 'settings':
        render_settings_view()

if __name__ == "__main__":
    main()