import streamlit as st
from transformers import pipeline

# Cache model to avoid reloading
@st.cache_resource
def load_model():
    return pipeline(
        "text2text-generation",
        model="google/flan-t5-small"   # smaller & safer for HF Spaces
    )

generator = load_model()

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

st.title("🏋️ Fitness Plan AI Generator")
st.write("Generate a personalized fitness plan using AI")

# User Inputs
age = st.number_input("Age", min_value=15, max_value=80, value=22)
gender = st.selectbox("Gender", ["Male", "Female"])
goal = st.selectbox("Fitness Goal", ["Weight Loss", "Muscle Gain", "General Fitness"])
level = st.selectbox("Fitness Level", ["Beginner", "Intermediate", "Advanced"])
days = st.slider("Workout Days per Week", 1, 7, 4)

if st.button("Generate Fitness Plan"):
    prompt = (
        f"Create a {days}-day fitness workout plan for a "
        f"{age}-year-old {gender}. "
        f"Goal: {goal}. "
        f"Fitness level: {level}. "
        f"Include warm-up, main workout, and cool-down."
    )

    with st.spinner("Generating your fitness plan..."):
        result = generator(prompt, max_length=300, do_sample=True)
        st.success("✅ Your Fitness Plan")
        st.write(result[0]["generated_text"])