import streamlit as st
from transformers import pipeline

# ---------------------------------------------------
# PAGE CONFIG
# ---------------------------------------------------
st.set_page_config(page_title="FitPlan AI", layout="centered")

st.title("🏋️ FitPlan AI – Personalized Fitness Plan")

# ---------------------------------------------------
# LOAD MODEL (Cached)
# ---------------------------------------------------
@st.cache_resource
def load_model():
    return pipeline(
        "text-generation",
        model="google/flan-t5-base"
    )

generator = load_model()

# ---------------------------------------------------
# USER NAME
# ---------------------------------------------------
name = st.text_input("Enter Your Name")

# ---------------------------------------------------
# GOALS
# ---------------------------------------------------
goal = st.selectbox(
    "Fitness Goal",
    [
        "Flexible",
        "Weight Loss",
        "Build Muscle",
        "Strength Gaining",
        "Abs Building"
    ]
)

# ---------------------------------------------------
# EQUIPMENT
# ---------------------------------------------------
st.subheader("Available Equipment")

col1, col2 = st.columns(2)

with col1:
    dumbbells = st.checkbox("Dumbbells")
    resistance_band = st.checkbox("Resistance Band")
    yoga_mat = st.checkbox("Yoga Mat")
    no_equipment = st.checkbox("No Equipment")
    inclined_bench = st.checkbox("Inclined Bench")
    treadmill = st.checkbox("Treadmill")
    cycle = st.checkbox("Cycle")

with col2:
    skipping_rope = st.checkbox("Skipping Rope")
    hand_gripper = st.checkbox("Hand Gripper")
    pullups_bar = st.checkbox("Pullups Bar")
    weight_plates = st.checkbox("Weight Plates")
    hula_hoop = st.checkbox("Hula Hoop Ring")
    bosu_ball = st.checkbox("Bosu Ball")

equipment = []

equipment_map = {
    "Dumbbells": dumbbells,
    "Resistance Band": resistance_band,
    "Yoga Mat": yoga_mat,
    "No Equipment": no_equipment,
    "Inclined Bench": inclined_bench,
    "Treadmill": treadmill,
    "Cycle": cycle,
    "Skipping Rope": skipping_rope,
    "Hand Gripper": hand_gripper,
    "Pullups Bar": pullups_bar,
    "Weight Plates": weight_plates,
    "Hula Hoop Ring": hula_hoop,
    "Bosu Ball": bosu_ball,
}

for item, selected in equipment_map.items():
    if selected:
        equipment.append(item)

# ---------------------------------------------------
# FITNESS LEVEL
# ---------------------------------------------------
st.subheader("Fitness Level")

fitness_level = st.radio(
    "",
    ["Beginner", "Intermediate", "Advanced"],
    horizontal=True
)

# ---------------------------------------------------
# PROMPT BUILDER
# ---------------------------------------------------
def build_prompt(name, goal, fitness_level, equipment):
    equipment_text = ", ".join(equipment)

    prompt = f"""
    Generate a detailed 5-day workout plan for {name}.

    Goal: {goal}
    Fitness Level: {fitness_level}
    Available Equipment: {equipment_text}

    Include:
    - Day-wise workout split
    - Exercises
    - Sets and reps
    - Rest time
    - Warm-up and cool-down
    """

    return prompt


# ---------------------------------------------------
# SUBMIT BUTTON
# ---------------------------------------------------
if st.button("Generate Personalised Plan"):

    if not name:
        st.error("Please enter your name.")
    elif not equipment:
        st.error("Please select at least one equipment option.")
    else:
        with st.spinner("Generating your AI workout plan..."):
            prompt = build_prompt(name, goal, fitness_level, equipment)

            result = generator(
                prompt,
                max_length=500,
                do_sample=True,
                temperature=0.7
            )

            plan = result[0]["generated_text"]

        st.success(f"Plan Generated Successfully, {name}! 💪")
        st.write(plan)