import streamlit as st
from Embeddings import GetEmbeddings
import json


# Load Agent once and cache it
@st.cache_resource
def load_agent():
    agent = GetEmbeddings(config_path="config.json")
    agent.run()                # Build/load FAISS
    agent.load_summarizer()    # Load summarizer model
    encoder = agent.load_encoder()
    return agent, encoder


def main():
    st.set_page_config(page_title="📊 Financial QA Agent", layout="wide")

    st.title("📊 Financial QA Agent")
    st.markdown(
        """
        Ask questions about financial reports.  
        The system retrieves relevant sections from company reports and summarizes them into concise answers.
        """
    )

    # Sidebar
    st.sidebar.header("⚙️ Settings")
    show_debug = st.sidebar.checkbox("Show retrieved chunks", value=False)

    # Load Agent
    agent, encoder = load_agent()

    # User Input
    query = st.text_area("Enter your financial question:", height=100)

    if st.button("Get Answer"):
        if query.strip() == "":
            st.warning("⚠️ Please enter a query.")
        else:
            with st.spinner("🔎 Searching and generating answer..."):
                answer = agent.answer_query(query, top_k=3)

            st.subheader("✅ Answer")
            st.write(answer)

            if show_debug:
                st.subheader("📂 Retrieved Chunks (Debug)")
                # Show top chunks used
                q_emb = encoder.encode(query, convert_to_numpy=True).reshape(1, -1)
                import faiss
                faiss.normalize_L2(q_emb)
                scores, idxs = agent.index.search(q_emb, k=3)
                for score, idx in zip(scores[0], idxs[0]):
                    st.markdown(f"**Score:** {score:.4f}")
                    st.write(agent.metadata[idx]["text"][:500] + "...")


if __name__ == "__main__":
    main()