import streamlit as st
import pandas as pd
from rag_agent import RAGAgent
from dashboard import DashboardBuilder
from vectorstore import VectorStore

st.set_page_config(page_title="📊 RAG AI Forecasting Agent", layout="wide")
st.title("🧠 Ask Your Data (RAG Powered)")

vs = VectorStore("vector_db")
rag_agent = RAGAgent(vs)
dashboard = DashboardBuilder()

# Upload multiple files
uploaded_files = st.file_uploader("Upload CSVs", type="csv", accept_multiple_files=True)
if uploaded_files:
    for file in uploaded_files:
        try:
            df = pd.read_csv(file)
            vs.add_dataframe(df, file.name)
            st.success(f"✅ {file.name} uploaded and indexed")
        except Exception as e:
            st.error(f"❌ Error uploading {file.name}: {str(e)}")

# Chat box
user_input = st.chat_input("Ask me anything about your data")
if user_input:
    st.chat_message("user").write(user_input)
    response = rag_agent.answer(user_input)
    st.chat_message("assistant").write(response)

    # Dashboard preview trigger
    if "dashboard" in user_input.lower():
        all_dfs = vs.get_all_dataframes()
        for name, df in all_dfs.items():
            st.subheader(f"🖋 Dashboard for {name}")
            dashboard.display(df)

    # Example of encoding user input and calculating similarity
    sentences = [user_input]  # You can add more sentences if needed
    embeddings = vs.encode_sentences(sentences)
    similarities = vs.calculate_similarity(embeddings)
    st.write("Similarity Matrix:", similarities)