import streamlit as st
import pandas as pd
from pymongo import MongoClient
from db import insert_data_if_empty, get_mongo_client  # Import functions from db.py
import google.generativeai as genai  # Import Generative AI library

#### **1. Ensure Data is Inserted Before Display**
insert_data_if_empty()

#### **2. MongoDB Connection**
collection = get_mongo_client()

#### **3. Streamlit App to Display Data**
st.title("📊 MongoDB Data Viewer with AI Chatbot")

# Show first 5 rows from MongoDB
st.subheader("First 5 Rows from Database")
data = list(collection.find({}, {"_id": 0}).limit(5))

if data:
    st.write(pd.DataFrame(data))
else:
    st.warning("⚠️ No data found. Try refreshing the app.")

# Button to show full MongoDB data
if st.button("Show Complete Data"):
    all_data = list(collection.find({}, {"_id": 0}))
    st.write(pd.DataFrame(all_data))

#### **4. GenAI Chatbot Interface**
st.subheader("🤖 AI Chatbot")

# Configure Generative AI model (replace `your-api-key` with actual key)
genai.configure(api_key="your-api-key")  # Set up your API key

# Initialize model
model = genai.GenerativeModel("gemini-pro")

# User input for chatbot
user_prompt = st.text_input("Ask AI something:")

if st.button("Get AI Response"):
    if user_prompt:
        response = model.generate_content(user_prompt)
        st.write("### AI Response:")
        st.write(response.text)
    else:
        st.warning("⚠️ Please enter a question.")