import streamlit as st
import os
from src.rag_engine import ProjectRAGEngine

# -------------------------------
# Page configuration
# -------------------------------
st.set_page_config(
    page_title="IIR Project Analyzer",
    layout="wide"
)

st.title("📂 Industrial Project Report Analyzer")

# -------------------------------
# Session persistence
# -------------------------------
if "engine" not in st.session_state:
    st.session_state.engine = ProjectRAGEngine()

if "processed_files" not in st.session_state:
    st.session_state.processed_files = []

if "messages" not in st.session_state:
    st.session_state.messages = []

# -------------------------------
# File upload
# -------------------------------
uploaded_files = st.sidebar.file_uploader(
    "Upload Project PDFs",
    type="pdf",
    accept_multiple_files=True
)

if uploaded_files:
    uploaded_names = {f.name for f in uploaded_files}

    if set(st.session_state.processed_files) != uploaded_names:
        with st.spinner("Analyzing project reports..."):
            temp_dir = "temp"
            os.makedirs(temp_dir, exist_ok=True)

            paths = []
            for f in uploaded_files:
                path = os.path.join(temp_dir, f.name)
                with open(path, "wb") as out:
                    out.write(f.getbuffer())
                paths.append(path)

            st.session_state.engine.process_documents(paths)
            st.session_state.processed_files = list(uploaded_names)

        st.success("Reports indexed. Ready for queries.")

# -------------------------------
# Display previous chat messages
# -------------------------------
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.write(msg["content"])

# -------------------------------
# Chat interface
# -------------------------------
query = st.chat_input("Ex: 'Compare the budgets of these projects'")

if query:
    # Store user message
    st.session_state.messages.append(
        {"role": "user", "content": query}
    )

    with st.chat_message("user"):
        st.write(query)

    with st.chat_message("assistant"):
        answer, sources = st.session_state.engine.get_answer(query)

        st.markdown("### Response")
        st.write(answer)

        # Store assistant message
        st.session_state.messages.append(
            {"role": "assistant", "content": answer}
        )

        if sources:
            with st.expander("📌 Source Attribution & Quotes"):
                for idx, s in enumerate(sources):
                    doc_name = os.path.basename(s["metadata"]["source"])
                    page_num = s["metadata"]["page"] + 1

                    st.markdown(
                        f"**Source {idx + 1}:** {doc_name} (Page {page_num})"
                    )
                    st.caption(f"\"{s['content'][:300]}...\"")