src/streamlit_app.py

import streamlit as st
import requests
import os

# Backend URL (change if deployed)
BACKEND_URL = "http://localhost:8000"

st.set_page_config(page_title="PDF Assistant", page_icon="📘", layout="wide")

# ---------------- CSS (Dark Theme) ----------------
# FIX: Added CSS for the footer
st.markdown("""
<style>
/* Streamlit standard setup for dark theme adherence */
:root {
    --primary-color: #1e3a8a; /* Blue for highlights */
    --background-color: #0e1117;
    --secondary-background-color: #1a1d29;
    --text-color: #f0f2f6;
}

/* Custom Chat Bubbles */
.chat-user {
    background: #2d3748; /* Dark gray */
    padding: 12px;
    border-radius: 10px 10px 2px 10px; /* Rounded corners for chat bubble */
    margin: 6px 0 6px auto;
    max-width: 85%;
    text-align: right;
    color: var(--text-color);
}
.chat-bot {
    background: var(--primary-color); /* Primary blue */
    padding: 12px;
    border-radius: 10px 10px 10px 2px;
    margin: 6px auto 6px 0;
    max-width: 85%;
    text-align: left;
    color: #ffffff; /* White text for contrast */
}

/* Sources section styling */
.sources {
    font-size: 0.8em;
    opacity: 0.7;
    margin-top: 10px;
    border-top: 1px solid rgba(255, 255, 255, 0.1);
    padding-top: 5px;
}

/* Footer styling */
.footer {
    position: fixed;
    left: 0;
    bottom: 0;
    width: 100%;
    background-color: var(--secondary-background-color);
    color: var(--text-color);
    text-align: center;
    padding: 10px;
    font-size: 0.85em;
    border-top: 1px solid rgba(255, 255, 255, 0.1);
}
.footer a {
    color: var(--primary-color);
    text-decoration: none;
    font-weight: bold;
}
.footer a:hover {
    text-decoration: underline;
}
</style>
""", unsafe_allow_html=True)

# ---------------- SESSION STATE ----------------
if "chat" not in st.session_state:
    st.session_state.chat = []

if "uploaded_file_name" not in st.session_state:
    st.session_state.uploaded_file_name = None

# Add a key to the file uploader to allow it to be reset.
if "uploader_key" not in st.session_state:
    st.session_state.uploader_key = 0

# FIX 1: Change application name
st.title("📘 PDF Assistant")


# ---------------- FUNCTIONS ----------------
def clear_chat_history():
    """Clears the chat history in the session state."""
    st.session_state.chat = []


def clear_memory():
    """Calls the backend endpoint to clear loaded PDF data and resets UI state."""
    res = requests.post(f"{BACKEND_URL}/clear")
    if res.status_code == 200:
        st.session_state.uploaded_file_name = None
        # Increment the key of the file uploader to clear its value
        st.session_state.uploader_key += 1
        st.success("Memory cleared. Please upload a new PDF.")
    else:
        st.error(f"Failed to clear memory: {res.json().get('detail', 'Unknown error')}")
    # Removed st.rerun() to prevent "no-op" warning


# ---------------- SIDEBAR CONTROLS ----------------
with st.sidebar:
    st.header("Controls")
    st.button("🗑️ Clear Chat History", on_click=clear_chat_history, use_container_width=True)
    st.button("🔥 Clear PDF Memory", on_click=clear_memory, use_container_width=True)

    st.markdown("---")
    if st.session_state.uploaded_file_name:
        st.success(f"✅ **Active PDF:**\n `{st.session_state.uploaded_file_name}`")
    else:
        st.warning("⬆️ Upload a PDF to start chatting!")

# ---------------- UPLOAD PDF ----------------
# Use the dynamic key for the file uploader.
uploaded = st.file_uploader(
    "Upload your PDF",
    type=["pdf"],
    key=st.session_state.uploader_key # Use the dynamic key
)

# Only process if a file is uploaded AND it's a NEW file
if uploaded and uploaded.name != st.session_state.uploaded_file_name:
    st.session_state.uploaded_file_name = None  # Clear status while processing
    st.session_state.chat = []  # Clear chat for a new document

    with st.spinner(f"Processing '{uploaded.name}'..."):
        try:
            files = {"file": (uploaded.name, uploaded.getvalue(), "application/pdf")}
            res = requests.post(f"{BACKEND_URL}/upload", files=files)

            if res.status_code == 200:
                chunks = res.json().get("chunks", 0)
                st.success(f"PDF processed successfully! {chunks} chunks created.")
                st.session_state.uploaded_file_name = uploaded.name
            else:
                error_msg = res.json().get("detail", "Unknown error during processing.")
                st.error(f"Upload failed: {error_msg}")
                st.session_state.uploaded_file_name = None

        except requests.exceptions.ConnectionError:
            st.error(f"Could not connect to the backend server at {BACKEND_URL}. Ensure it is running.")
            st.session_state.uploaded_file_name = None
        except Exception as e:
            st.error(f"An unexpected error occurred: {e}")
            st.session_state.uploaded_file_name = None

    # Rerun the app to update the UI status immediately
    st.rerun()

# ---------------- CHAT INPUT ----------------
# Disable input field if no PDF is loaded
disabled_input = st.session_state.uploaded_file_name is None
question = st.text_input(
    "Ask a question about the loaded PDF:",
    key="question_input",
    disabled=disabled_input
)

if st.button("Send", disabled=disabled_input) and question:
    # 1. Add user query to chat history
    st.session_state.chat.append(("user", question))

    # 2. Call backend
    with st.spinner("Thinking..."):
        try:
            res = requests.post(f"{BACKEND_URL}/ask", json={"question": question})

            if res.status_code == 200:
                data = res.json()
                answer = data.get("answer", "No answer provided.")
                sources = data.get("sources", 0)

                # Format the bot's response to include source count
                bot_message = f"{answer}<div class='sources'>Context Chunks Used: {sources}</div>"
                st.session_state.chat.append(("bot", bot_message))

            else:
                error_detail = res.json().get("detail", "Error while generating answer.")
                st.session_state.chat.append(("bot", f"🔴 **Error:** {error_detail}"))

        except requests.exceptions.ConnectionError:
            st.session_state.chat.append(("bot",
                                          f"🔴 **Error:** Could not connect to the backend server at {BACKEND_URL}. Ensure it is running."))
        except Exception as e:
            st.session_state.chat.append(("bot", f"🔴 **An unexpected error occurred:** {e}"))

    # Rerun to display the updated chat history
    st.rerun()

# ---------------- SHOW CHAT HISTORY ----------------
st.markdown("## Chat History")
# Reverse the list to show the latest messages at the bottom
for role, msg in st.session_state.chat:
    if role == "user":
        st.markdown(f"<div class='chat-user'>{msg}</div>", unsafe_allow_html=True)
    else:
        # Bot message includes the source count, so use the HTML content
        st.markdown(f"<div class='chat-bot'>{msg}</div>", unsafe_allow_html=True)

# ---------------- FOOTER (Creator Credit) ----------------
# FIX 2: Add creator credit with LinkedIn link
footer_html = """
<div class="footer">
    Created by <a href="https://www.linkedin.com/in/abhishek-iitr/" target="_blank">Abhishek Saxena</a>
</div>
"""
st.markdown(footer_html, unsafe_allow_html=True)

# ---------------- END ----------------