import streamlit as st
from auth import generate_otp, create_jwt, verify_jwt
from email_utils import send_otp_email  
from dotenv import load_dotenv
load_dotenv()

st.set_page_config(page_title="FitPlan AI Login")

# Initialize session state
if "otp" not in st.session_state:
    st.session_state.otp = None

if "authenticated" not in st.session_state:
    st.session_state.authenticated = False

if "token" not in st.session_state:
    st.session_state.token = None


st.title("🔐 FitPlan AI - OTP Login")

# ------------------------
# STEP 1: Enter Email
# ------------------------
if not st.session_state.authenticated:

    email = st.text_input("Enter your Email")

    if st.button("Send OTP"):
        if email:
            otp = generate_otp()
            st.session_state.otp = otp
            st.session_state.email = email
            
            send_otp_email(email, otp)
            
            st.success("OTP sent to your email 📧")
        else:
            st.error("Please enter email")

    # ------------------------
    # STEP 2: Verify OTP
    # ------------------------
    if st.session_state.otp:
        entered_otp = st.text_input("Enter OTP")

        if st.button("Verify OTP"):
            if entered_otp == st.session_state.otp:
                token = create_jwt(st.session_state.email)
                st.session_state.token = token
                st.session_state.authenticated = True
                st.success("Login Successful!")
                st.rerun()
            else:
                st.error("Invalid OTP")


# ------------------------
# Protected Page
# ------------------------
else:
    decoded = verify_jwt(st.session_state.token)

    if decoded:
        st.success(f"Welcome {decoded['email']} 🎉")

        st.write("You can now access FitPlan AI")

        if st.button("Logout"):
            st.session_state.authenticated = False
            st.session_state.token = None
            st.session_state.otp = None
            st.rerun()
    else:
        st.error("Session Expired")