import streamlit as st

st.set_page_config(page_title="FitPlan AI", layout="centered")

# ---------------- SESSION STATE ----------------
if "step" not in st.session_state:
    st.session_state.step = 1

# ---------------- STEP 1: LANDING PAGE ----------------
if st.session_state.step == 1:
    st.title("💪 FitPlan AI")
    st.subheader("Your Personalized Fitness Plan Generator")

    st.markdown("""
    Welcome to **FitPlan AI** 🚀  
    Get a customized workout plan based on:
    - Your fitness profile
    - Your goals
    - Available equipment
    """)

    if st.button("Generate My Plan"):
        st.session_state.step = 2

# ---------------- STEP 2: FITNESS PROFILE ----------------
elif st.session_state.step == 2:
    st.title("👤 Your Fitness Profile")

    age = st.number_input("Age", min_value=10, max_value=80)
    gender = st.selectbox("Gender", ["Male", "Female", "Other"])
    weight = st.number_input("Weight (kg)")
    height = st.number_input("Height (cm)")
    experience = st.selectbox(
        "Fitness Level",
        ["Beginner", "Intermediate", "Advanced"]
    )

    if st.button("Next"):
        st.session_state.profile = {
            "age": age,
            "gender": gender,
            "weight": weight,
            "height": height,
            "experience": experience
        }
        st.session_state.step = 3

# ---------------- STEP 3: SELECT GOAL ----------------
elif st.session_state.step == 3:
    st.title("🎯 Select Your Goal")

    goal = st.radio(
        "Choose your fitness goal:",
        ["Increase Flexibility 🧘", 
         "Build Muscle 💪", 
         "Weight Loss / Get Lean 🔥"]
    )

    if st.button("Next"):
        st.session_state.goal = goal
        st.session_state.step = 4

# ---------------- STEP 4: SELECT EQUIPMENT ----------------
elif st.session_state.step == 4:
    st.title("🏋️ Select Available Equipment")

    equipment = st.multiselect(
        "Select equipment you have access to:",
        [
            "Dumbbells",
            "Barbell",
            "Resistance Bands",
            "Treadmill",
            "Stationary Bike",
            "Bench Press",
            "Pull-up Bar",
            "Kettlebells",
            "Leg Press Machine",
            "Cable Machine",
            "Smith Machine",
            "No Equipment (Bodyweight Only)"
        ]
    )

    if st.button("Generate Plan"):
        st.session_state.equipment = equipment
        st.session_state.step = 5

# ---------------- STEP 5: DISPLAY PLAN ----------------
elif st.session_state.step == 5:
    st.title("📄 Your Personalized Fitness Plan")

    st.subheader("Profile Summary")
    st.write(st.session_state.profile)
    st.write("Goal:", st.session_state.goal)
    st.write("Equipment:", st.session_state.equipment)

    st.divider()

    # Simple frontend demo output
    if "Build Muscle" in st.session_state.goal:
        st.success("💪 Suggested Plan: 4-Day Strength Split")
        st.markdown("""
        - Day 1: Chest + Triceps  
        - Day 2: Back + Biceps  
        - Day 3: Legs  
        - Day 4: Shoulders + Core  
        """)

    elif "Flexibility" in st.session_state.goal:
        st.success("🧘 Suggested Plan: 5-Day Mobility Routine")
        st.markdown("""
        - Dynamic Stretching  
        - Yoga Flow  
        - Hamstring & Hip Mobility  
        - Shoulder Mobility  
        - Core Stability  
        """)

    else:
        st.success("🔥 Suggested Plan: 4-Day Fat Loss Program")
        st.markdown("""
        - HIIT Training  
        - Strength + Cardio Mix  
        - Core Workouts  
        - 30 min Daily Cardio  
        """)

    if st.button("Start Over"):
        st.session_state.step = 1