import streamlit as st
import streamlit.components.v1 as components
from huggingface_hub import InferenceClient
import os


# ================= BMI FUNCTIONS =================

def calculate_bmi(weight, height):
    if height == 0:
        return 0
    height_m = height / 100
    return weight / (height_m ** 2)



def bmi_category(bmi):
    if bmi < 18.5:
        return "Underweight"
    elif bmi < 25:
        return "Normal Weight"
    elif bmi < 30:
        return "Overweight"
    else:
        return "Obese"


def build_prompt(name, gender, height, weight, goal, fitness_level, equipment):

    bmi = calculate_bmi(weight, height)
    bmi_status = bmi_category(bmi)

    equipment_list = ", ".join(equipment) if equipment else "No Equipment"

    prompt = f"""
You are a certified professional fitness trainer.

Create a structured 5-day personalized workout plan.

User Profile:
- Name: {name}
- Gender: {gender}
- Height: {height} cm
- Weight: {weight} kg
- BMI: {bmi:.2f} ({bmi_status})
- Goal: {goal}
- Fitness Level: {fitness_level}
- Available Equipment: {equipment_list}

Instructions:
1. Divide clearly into Day 1 to Day 5.
2. Include exercise name.
3. Include sets and reps.
4. Include rest period.
5. Adjust intensity based on BMI category.
6. Avoid unsafe exercises for beginners.
7. Keep the plan professional and easy to follow.
"""

    return prompt, bmi, bmi_status



# ================= QUERY MODEL =================

def query_model(prompt):

    try:

        HF_TOKEN = os.getenv("HF_TOKENN")

        client = InferenceClient(
            model="mistralai/Mistral-7B-Instruct-v0.2",
            token=HF_TOKEN
        )

        response = client.chat_completion(

            messages=[
                {"role": "system", "content": "You are a certified professional fitness trainer."},
                {"role": "user", "content": prompt}
            ],

            max_tokens=600,
            temperature=0.7
        )

        return response.choices[0].message.content


    except Exception as e:

        return f"Error: {str(e)}"



# ================= PAGE CONFIG =================

st.set_page_config(
    page_title="FitPlan AI",
    page_icon="💪",
    layout="wide"
)



# ================= SIDEBAR =================

st.sidebar.title("💪 FitPlan AI")

page = st.sidebar.radio(
    "Navigate",
    ["🏠 Fitness Form", "📊 BMI Result"]
)



# ================= FORM PAGE =================

if page == "🏠 Fitness Form":

    st.title("🏠 Fitness Profile Form")

    name = st.text_input("Name")

    gender = st.selectbox(
        "Gender",
        ["Male", "Female"]
    )

    fitness_level = st.selectbox(
        "Fitness Level",
        ["Beginner", "Intermediate", "Advanced"]
    )

    height_cm = st.number_input(
        "Height (cm)",
        min_value=0.0
    )

    weight_kg = st.number_input(
        "Weight (kg)",
        min_value=0.0
    )

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

    equipment = st.multiselect(
        "Available Equipment",
        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment"]
    )

    submit = st.button("Submit Profile")

    if submit:

        bmi = calculate_bmi(weight_kg, height_cm)

        category = bmi_category(bmi)

        st.session_state["name"] = name
        st.session_state["gender"] = gender
        st.session_state["fitness_level"] = fitness_level
        st.session_state["bmi"] = bmi
        st.session_state["category"] = category
        st.session_state["goal"] = goal
        st.session_state["equipment"] = equipment

        # ADD THESE TWO (IMPORTANT)
        st.session_state["height"] = height_cm
        st.session_state["weight"] = weight_kg


        st.success("Profile submitted successfully! Go to BMI Result page.")



# ================= RESULT PAGE =================

elif page == "📊 BMI Result":

    st.title("📊 BMI Result")

    if "bmi" not in st.session_state:

        st.warning("Please fill the Fitness Form first.")

    else:

        st.write(f"Name: {st.session_state['name']}")
        st.write(f"Gender: {st.session_state['gender']}")
        st.write(f"Fitness Level: {st.session_state['fitness_level']}")
        st.write(f"BMI: {st.session_state['bmi']}")
        st.write(f"Category: {st.session_state['category']}")
        st.write(f"Goal: {st.session_state['goal']}")
        st.write(f"Equipment: {', '.join(st.session_state['equipment'])}")



        gender = st.session_state["gender"]
        bmi_status = st.session_state["category"]
        goal = st.session_state["goal"]
        fitness_level = st.session_state["fitness_level"]
        equipment_list = ", ".join(st.session_state["equipment"])



        prompt, bmi, bmi_status = build_prompt(

    st.session_state["name"],

    st.session_state["gender"],

    st.session_state["height"],

    st.session_state["weight"],

    st.session_state["goal"],

    st.session_state["fitness_level"],

    st.session_state["equipment"]

)




        # FIXED POSITION
        with st.spinner("Generating AI Fitness Plan..."):

            result = query_model(prompt)



        st.subheader("AI Fitness Plan")

        st.write(result)