Update app.py

app.py

@@ -71,11 +71,20 @@ st.markdown("""
         transform: translateY(-2px) rotateX(15deg);
         box-shadow: 0 5px 15px rgba(0,0,0,0.3);
     }
+    
+    /* Admin badge */
+    .admin-badge {
+        background: #4a148c;
+        color: white !important;
+        padding: 2px 8px;
+        border-radius: 12px;
+        font-size: 0.8em;
+    }
 </style>
 """, unsafe_allow_html=True)
 
 import firebase_admin
-from firebase_admin import credentials, firestore, auth
+from firebase_admin import credentials, firestore, auth, exceptions
 import pandas as pd
 import plotly.express as px
 from datetime import datetime, timedelta
@@ -87,11 +96,40 @@ if not firebase_admin._apps:
     firebase_admin.initialize_app(cred)
 db = firestore.client()
 
+# Create default admin user if not exists
+def create_default_admin():
+    try:
+        admin_email = "admin@example.com"
+        try:
+            admin_user = auth.get_user_by_email(admin_email)
+        except auth.UserNotFoundError:
+            admin_user = auth.create_user(
+                email=admin_email,
+                password="Temp@2025",
+                display_name="Admin"
+            )
+            auth.set_custom_user_claims(admin_user.uid, {'admin': True})
+            db.collection("users").document(admin_user.uid).set({
+                "email": admin_email,
+                "is_admin": True,
+                "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            })
+        except Exception as e:
+            print(f"Error creating admin user: {str(e)}")
+    except Exception as e:
+        print(f"Error in admin setup: {str(e)}")
+
+create_default_admin()
+
 # Authentication Functions
 def register_user(email, password):
     try:
         user = auth.create_user(email=email, password=password)
-        db.collection("users").document(user.uid).set({"email": email})
+        db.collection("users").document(user.uid).set({
+            "email": email,
+            "is_admin": False,
+            "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        })
         st.success("Registration successful! Please sign in.")
     except Exception as e:
         st.error(f"Error: {e}")
@@ -99,10 +137,10 @@ def register_user(email, password):
 def authenticate_user(email, password):
     try:
         user = auth.get_user_by_email(email)
-        return user is not None
+        return user
     except Exception as e:
         st.error(f"Authentication error: {str(e)}")
-        return False
+        return None
 
 # Main Application Logic
 def main():
@@ -110,6 +148,7 @@ def main():
     if "authenticated" not in st.session_state:
         st.session_state.authenticated = False
         st.session_state.email = ""
+        st.session_state.is_admin = False
         st.session_state.first_login = True
         st.session_state.user_uid = ""
 
@@ -135,12 +174,15 @@ def main():
                 password = st.text_input("Password", type="password")
                 if st.form_submit_button("Login"):
                     try:
-                        if authenticate_user(email, password):
+                        user = authenticate_user(email, password)
+                        if user:
+                            custom_claims = user.custom_claims or {}
+                            user_doc = db.collection("users").document(user.uid).get().to_dict()
+                            
                             st.session_state.authenticated = True
-                            st.session_state.email = email
-                            # Get user UID for password updates
-                            user = auth.get_user_by_email(email)
+                            st.session_state.email = user.email
                             st.session_state.user_uid = user.uid
+                            st.session_state.is_admin = custom_claims.get('admin', False) or user_doc.get('is_admin', False)
                             st.rerun()
                         else:
                             st.error("Invalid credentials")
@@ -199,15 +241,23 @@ def main():
         st.session_state.first_login = False
 
     # Sidebar Navigation
-    st.sidebar.title(f"Welcome, {st.session_state.email}")
-    menu = st.sidebar.radio("Navigation", [
+    menu_items = [
         "🏠 Dashboard",
         "📝 Task Entry",
         "👀 View Tasks",
         "✏️ Edit Tasks",
-        "🏗️ Projects",
-        "⚙️ Settings"
-    ], label_visibility="collapsed")
+    ]
+    
+    if st.session_state.is_admin:
+        menu_items += ["🏗️ Projects"]
+    
+    menu_items += ["⚙️ Settings"]
+    
+    st.sidebar.title(f"Welcome, {st.session_state.email}")
+    if st.session_state.is_admin:
+        st.sidebar.markdown('<span class="admin-badge">ADMIN</span>', unsafe_allow_html=True)
+    
+    menu = st.sidebar.radio("Navigation", menu_items, label_visibility="collapsed")
 
     # Dashboard View
     if menu == "🏠 Dashboard":
@@ -289,7 +339,7 @@ def main():
                 except Exception as e:
                     st.error(f"Error saving task: {str(e)}")
 
-    # View Tasks with Filtering
+    # View Tasks
     elif menu == "👀 View Tasks":
         st.title("🔍 Task Explorer")
         tasks = get_tasks()
@@ -366,8 +416,12 @@ def main():
         else:
             st.info("No tasks available for editing")
 
-    # Project Management
+    # Project Management (Admin Only)
     elif menu == "🏗️ Projects":
+        if not st.session_state.is_admin:
+            st.error("⛔ You don't have permission to access this section")
+            return
+        
         st.title("🏗️ Project Management")
         try:
             projects_ref = db.collection("projects")