import streamlit as st
import firebase_admin
from firebase_admin import credentials, db
from transformers import RagRetriever, RagTokenizer, RagTokenForGeneration
from datasets import load_dataset
import torch
import json
from datasets import load_dataset

# Loading the dataset with trust_remote_code=True
dataset = load_dataset("wiki_dpr", "psgs_w100.nq.exact", trust_remote_code=True)

import os
import firebase_admin
from firebase_admin import credentials, initialize_app, db

# Get Firebase credentials and database URL from environment variables
firebase_credentials = json.loads(os.getenv("FIREBASE_CREDENTIALS"))
database_url = os.getenv("FIREBASE_DATABASE_URL")  # This is the new URL variable

if firebase_credentials is None or database_url is None:
    raise ValueError("FIREBASE_CREDENTIALS or FIREBASE_DATABASE_URL environment variables are not set.")

# Initialize Firebase Admin SDK with credentials and database URL
cred = credentials.Certificate(firebase_credentials)
initialize_app(cred, {'databaseURL': database_url})

# Example code to test database connection (writing and reading from Firebase)
ref = db.reference('/test')
ref.set({
    'name': 'Test Value'
})
print(ref.get())


# Firebase Database Reference
firebase_ref = db.reference('/projects')

# RAG Model Initialization
@st.cache_resource
def initialize_rag_model():
    retriever = RagRetriever.from_pretrained("facebook/rag-token-base")
    tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
    model = RagTokenForGeneration.from_pretrained("facebook/rag-token-base")
    return model, tokenizer, retriever

rag_model, rag_tokenizer, rag_retriever = initialize_rag_model()

# Sidebar for Navigation
st.sidebar.title("Project Management Dashboard")
menu = st.sidebar.radio("Menu", ["Dashboard", "Live Updates", "Ask RAG"])

# Dashboard Section
if menu == "Dashboard":
    st.title("Project Management Dashboard")
    st.write("This dashboard provides an overview of your project statuses, deadlines, and tasks.")
    
    # Fetch project data from Firebase
    projects = firebase_ref.get()  # Fetch data from Firebase
    if projects:
        for project_id, project in projects.items():
            st.subheader(f"Project: {project.get('name', 'Unknown')}")
            st.write(f"**Deadline:** {project.get('deadline', 'No Deadline')}")
            st.write(f"**Status:** {project.get('status', 'No Status')}")
            st.write(f"**Tasks:**")
            for task in project.get('tasks', []):
                st.write(f"- {task}")
    else:
        st.write("No projects found in the database.")

# Live Updates Section
elif menu == "Live Updates":
    st.title("Live Project Updates")
    st.write("Monitor live updates from Firebase in real-time.")
    
    # Add a button to trigger a database update
    if st.button("Trigger Live Update"):
        new_update = {
            "name": "New Project",
            "deadline": "2025-01-15",
            "status": "In Progress",
            "tasks": ["Task 1", "Task 2", "Task 3"]
        }
        firebase_ref.push(new_update)  # Push data to Firebase
        st.success("New update pushed to Firebase!")

# Ask RAG Section
elif menu == "Ask RAG":
    st.title("Ask RAG")
    st.write("Ask questions related to project data using the RAG model.")
    
    question = st.text_input("Enter your question:")
    if question:
        inputs = rag_tokenizer(question, return_tensors="pt")
        with torch.no_grad():
            outputs = rag_model.generate(**inputs)
        answer = rag_tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.write(f"**Answer:** {answer}")