Update app.py

app.py

@@ -1,35 +1,45 @@
 import streamlit as st
 
-from database import create_tables, save_weight, get_weight_history
+from database import create_tables, save_weight, get_weight_history, save_otp
 from auth import signup_user, login_user, verify_user_otp, generate_otp
 from email_utils import send_otp_email
 from model_api import query_model
 from prompt_builder import build_prompt
-from database import create_tables
 
-# create tables automatically
+
+# Create database tables automatically
 create_tables()
 
 st.set_page_config(page_title="FitPlan AI", layout="centered")
 
 st.title("🏋️ FitPlan AI - Personalized Fitness Planner")
 
-# session state
+
+# -------------------------
+# SESSION STATE
+# -------------------------
+
 if "user" not in st.session_state:
     st.session_state.user = None
 
-if "otp_sent" not in st.session_state:
-    st.session_state.otp_sent = False
-
 if "email" not in st.session_state:
     st.session_state.email = None
 
+if "credentials_verified" not in st.session_state:
+    st.session_state.credentials_verified = False
+
+if "otp_sent" not in st.session_state:
+    st.session_state.otp_sent = False
+
 
 # -------------------------
-# LOGIN / SIGNUP PAGE
+# SIDEBAR MENU
 # -------------------------
 
-menu = st.sidebar.selectbox("Menu", ["Login", "Signup"])
+menu = st.sidebar.selectbox(
+    "Menu",
+    ["Signup", "Login"]
+)
 
 
 # -------------------------
@@ -38,7 +48,7 @@ menu = st.sidebar.selectbox("Menu", ["Login", "Signup"])
 
 if menu == "Signup":
 
-    st.subheader("Create Account")
+    st.header("Create Account")
 
     email = st.text_input("Email")
     password = st.text_input("Password", type="password")
@@ -59,43 +69,55 @@ if menu == "Signup":
 
 if menu == "Login" and st.session_state.user is None:
 
-    st.subheader("Login")
+    st.header("Login")
 
     email = st.text_input("Email")
     password = st.text_input("Password", type="password")
 
-    if st.button("Send OTP"):
+    # Step 1: Verify credentials
+    if st.button("Login"):
 
         user = login_user(email, password)
 
         if user:
 
-            otp = generate_otp()
-
-            send_otp_email(email, otp)
-
-            from database import save_otp
-            save_otp(email, otp)
-
-            st.session_state.otp_sent = True
             st.session_state.email = email
+            st.session_state.credentials_verified = True
 
-            st.success("OTP sent to your email")
+            st.success("Credentials verified. Click Send OTP.")
 
         else:
             st.error("Invalid email or password")
 
 
+# -------------------------
+# SEND OTP
+# -------------------------
+
+if st.session_state.credentials_verified and not st.session_state.otp_sent:
+
+    if st.button("Send OTP"):
+
+        otp = generate_otp()
+
+        save_otp(st.session_state.email, otp)
+
+        send_otp_email(st.session_state.email, otp)
+
+        st.session_state.otp_sent = True
+
+        st.success("OTP sent to your email")
+
 
 # -------------------------
-# OTP VERIFICATION
+# VERIFY OTP
 # -------------------------
 
 if st.session_state.otp_sent:
 
-    st.subheader("Verify OTP")
+    st.subheader("Enter OTP")
 
-    otp_input = st.text_input("Enter OTP")
+    otp_input = st.text_input("OTP")
 
     if st.button("Verify OTP"):
 
@@ -104,6 +126,7 @@ if st.session_state.otp_sent:
         if verified:
 
             st.session_state.user = st.session_state.email
+            st.session_state.credentials_verified = False
             st.session_state.otp_sent = False
 
             st.success("Login successful")
@@ -112,9 +135,8 @@ if st.session_state.otp_sent:
             st.error("Invalid OTP")
 
 
-
 # -------------------------
-# MAIN DASHBOARD
+# DASHBOARD
 # -------------------------
 
 if st.session_state.user:
@@ -126,18 +148,12 @@ if st.session_state.user:
         ["Dashboard", "Weight Tracker", "Generate Fitness Plan"]
     )
 
-
-
-    # -------------------------
-    # DASHBOARD
-    # -------------------------
-
+    # Dashboard
     if page == "Dashboard":
 
         st.header("Welcome to FitPlan AI")
 
-        st.write("Track your fitness journey and generate personalized plans.")
-
+        st.write("Track your fitness progress and generate AI fitness plans.")
 
 
     # -------------------------
@@ -163,8 +179,7 @@ if st.session_state.user:
         if history:
             st.table(history)
         else:
-            st.info("No weight history found")
-
+            st.info("No weight data available")
 
 
     # -------------------------
@@ -193,14 +208,13 @@ if st.session_state.user:
 
             prompt = build_prompt(age, height, weight, goal, activity)
 
-            with st.spinner("Generating plan..."):
+            with st.spinner("Generating AI plan..."):
 
-                response = query_model(prompt)
+                result = query_model(prompt)
 
             st.subheader("Your AI Fitness Plan")
 
-            st.write(response)
-
+            st.write(result)
 
 
     # -------------------------
@@ -211,6 +225,7 @@ if st.session_state.user:
 
         st.session_state.user = None
         st.session_state.email = None
+        st.session_state.credentials_verified = False
         st.session_state.otp_sent = False
 
-        st.experimental_rerun()
+        st.rerun()