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")

# -------------------------------
# Environment validation
# -------------------------------
# openai_key = os.getenv("OPENAI_API_KEY")
# openrouter_key = os.getenv("OPENROUTER_API_KEY")

# if not openai_key or not openrouter_key:
#     st.error(
#         "Missing API keys.\n\n"
#         "Please set the following environment variables:\n"
#         "- OPENAI_API_KEY (for embeddings)\n"
#         "- OPENROUTER_API_KEY (for LLM)"
#     )
#     st.stop()

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

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

if uploaded_files:
    if (
        "processed_files" not in st.session_state
        or set(st.session_state.processed_files) != {f.name for f in uploaded_files}
    ):
        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 = [f.name for f in uploaded_files]

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

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

if 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)

        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]}...\"")