import streamlit as st
import base64
import os
from rag_core import process_pdf_to_vectorstore, get_rag_answer
# from langchain_community.callbacks.streamlit import StreamlitCallbackHandler

# --- PFAD ZUM BILD ---
# Findet heraus, wo app.py liegt, gehe einen Ordner hoch und gehen in 'assets'
current_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(current_dir)
svg_path = os.path.join(root_dir, "assets", "pdfi-avatar.svg")
favicon_path = os.path.join(root_dir, "assets", "favicon-pdfi.svg")

# --- PAGE CONFIG ---
# Hier wird die Variable genutzt.
st.set_page_config(page_title="Ask PDFi", page_icon=favicon_path, layout="centered")

# --- CUSTOM CSS (Disclaimer unter Chat) ---
st.markdown("""
<style>
    /* Disclaimer direkt unter dem Streamlit Chat-Eingabefeld anheften */
    [data-testid="stChatInput"]::after {
        content: "Ask PDFi kann Fehler machen. Bitte überprüfe wichtige Informationen. Deine Daten werden nur lokal und temporär im RAM verarbeitet.";
        display: block;
        font-size: 0.75rem;
        color: #888888;
        text-align: center;
        margin-top: 8px;
    }
</style>
""", unsafe_allow_html=True)

# --- SESSION STATE ---
# Speichert Chat-Verlauf und Status, damit beim Neuladen nichts verschwindet
if "messages" not in st.session_state:
    st.session_state.messages = []
if "pdf_processed" not in st.session_state:
    st.session_state.pdf_processed = False

# --- SIDEBAR (Upload) ---
with st.sidebar:
    st.header("📄 Dokumenten-Upload")
    uploaded_file = st.file_uploader("PDF hier ablegen", type=["pdf"])
    
    if uploaded_file is None:
        st.session_state.pdf_processed = False
        st.info("Bitte lade ein PDF hoch, um zu starten.")
    else:
        if st.button("PDF verarbeiten"):
            # Lade-Text an die Eule angepasst
            with st.spinner("PDFi überfliegt das Dokument..."):
                # ECHTE FUNKTION AUFRUFEN UND SPEICHERN
                st.session_state.vectorstore = process_pdf_to_vectorstore(uploaded_file)
                st.session_state.pdf_processed = True
            st.success("Dokument erfolgreich im Nest abgelegt!")

# --- IMPRESSUM & DATENSCHUTZ (Ganz unten in der Sidebar) ---
    st.markdown("---") 
    with st.expander("⚖️ Impressum & Datenschutz"):
        st.markdown("""
        **Angaben gemäß § 5 TMG:**
        
        Lars Patzenbein
             
        **Kontakt:** E-Mail: kontakt@askpdfi.com
        
        ---
        **Datenschutz & Security:** Diese Anwendung ist ein lokaler RAG-Prototyp. Sie läuft vollständig autark. Alle hochgeladenen Dokumente und Chat-Verläufe werden **ausschließlich temporär im Arbeitsspeicher (RAM)** verarbeitet. 
        
        Es findet keine Speicherung auf Festplatten statt und es fließen **keine Daten an externe APIs oder Drittanbieter**. Nach Beendigung der Sitzung oder einem Neuladen der Seite werden alle Daten restlos gelöscht.
        """)

# --- MAIN CHAT AREA ---
# Bild sicher für HTML konvertieren
def get_image_base64(path):
    with open(path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode()

svg_base64 = get_image_base64(svg_path)

# Logo und Titel in einer Zeile mit etwas Abstand anzeigen
st.markdown(f"""
<div style="display: flex; align-items: center; margin-bottom: 15px;">
    <img src="data:image/svg+xml;base64,{svg_base64}" width="75" height="75" style="margin-right: 15px;">
    <h1 style="margin: 0; padding: 0;">Ask PDFi</h1>
</div>
""", unsafe_allow_html=True)

st.markdown("""Deine smarte Tech-Eule für Dokumentenanalyse. Stelle komplexe Fragen zu deinen PDFs und **PDFi** findet die Antworten mit der Power von **Llama 3.1 8B**.""") 
st.markdown("<hr style='border: 1px solid #FF6600; opacity: 0.5;'>", unsafe_allow_html=True)

# Bisherige Chat-Historie anzeigen
for message in st.session_state.messages:
    # Avatar Zuweisung: Eule für die KI, Standard-User-Symbol für User
    avatar = svg_path if message["role"] == "assistant" else "👤"
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

# --- CHAT INPUT ---
# Inputfeld deaktivieren, wenn kein PDF verarbeitet wurde
if not st.session_state.pdf_processed:
    st.chat_input("Bitte lade zuerst ein PDF in der Seitenleiste hoch.", disabled=True)
else:
    # Wenn der User eine Frage tippt
    if prompt := st.chat_input("Was möchtest du aus dem PDF wissen?"):
        
        # User Message anzeigen und im State speichern
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user", avatar="👤"):
            st.markdown(prompt)

        # KI Antwort mit dem echten Gehirn generieren
        with st.chat_message("assistant", avatar=svg_path):
            # Er fängt den Stream aus der rag_core ab und schreibt ihn live.
                        
            # ECHTE FUNKTION AUFRUFEN
            answer = st.write_stream(get_rag_answer(st.session_state.vectorstore, prompt))
            
        # Die Antwort der KI in der Historie speichern!
        st.session_state.messages.append({"role": "assistant", "content": answer})