import streamlit as st
from huggingface_hub import InferenceClient
import os

# -------------------------
# PAGE CONFIG
# -------------------------
st.set_page_config(page_title="FitPlan-AI", page_icon="💪")
st.title("💪 FitPlan-AI: Personalized Fitness Profile")

# -------------------------
# LOAD HF CLIENT (CACHED)
# -------------------------
@st.cache_resource
def get_hf_client():
    hf_token = os.getenv("HF_TOKEN")

    if not hf_token:
        raise ValueError("HF_TOKEN not set in environment variables.")

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

# -------------------------
# BMI FUNCTIONS
# -------------------------
def calculate_bmi(weight, height_cm):
    height_m = height_cm / 100
    return round(weight / (height_m ** 2), 2)

def get_category(bmi):
    if bmi < 18.5:
        return "Underweight"
    elif bmi < 24.9:
        return "Normal"
    elif bmi < 29.9:
        return "Overweight"
    else:
        return "Obese"

# -------------------------
# PROMPT BUILDER
# -------------------------
def build_prompt(name, height, weight, goal, level, equipment, bmi, bmi_status):
    equipment_list = ", ".join(equipment)

    prompt = f"""
Create a detailed 5-day personalized workout plan.

Client Details:
Name: {name}
Height: {height} cm
Weight: {weight} kg
BMI: {bmi} ({bmi_status})
Goal: {goal}
Fitness Level: {level}
Available Equipment: {equipment_list}

Instructions:
- Provide a structured 5-day plan.
- Each day must include:
  - Warm-up
  - Main Workout (Sets x Reps)
  - Rest time
  - Cooldown
- Ensure exercises match fitness level and equipment.
- Keep it practical and realistic.
- Format clearly as:

Day 1:
Warm-up:
Main Workout:
Cooldown:

Continue up to Day 5.
"""

    return prompt.strip()

# -------------------------
# MODEL QUERY
# -------------------------
def query_model(prompt):
    try:
        client = get_hf_client()

        response = client.chat_completion(
            messages=[
                {"role": "system", "content": "You are a certified professional fitness trainer."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=800,
            temperature=0.7,
            top_p=0.9
        )

        return response.choices[0].message.content.strip()

    except Exception as e:
        return f"⚠️ Model Error: {str(e)}"

# -------------------------
# FORM
# -------------------------
with st.form("fitness_form"):

    st.subheader("Personal Information")

    name = st.text_input("Full Name*", placeholder="Enter your name")

    col1, col2 = st.columns(2)
    with col1:
        height = st.number_input("Height (cm)*", min_value=1.0, step=0.1)
    with col2:
        weight = st.number_input("Weight (kg)*", min_value=1.0, step=0.1)

    st.subheader("Fitness Details")

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

    level = st.radio(
        "Fitness Level",
        ["Beginner", "Intermediate", "Advanced"]
    )

    equipment = st.multiselect(
        "Available Equipment",
        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment", 
         "Kettlebell", "Pull-up Bar"]
    )

    submit = st.form_submit_button("Submit Profile")

# -------------------------
# HANDLE SUBMISSION
# -------------------------
if submit:

    if not name:
        st.error("Please enter your name.")

    elif height <= 0 or weight <= 0:
        st.error("Please enter valid height and weight.")

    elif not equipment:
        st.error("Please select at least one equipment option.")

    else:
        st.success("Profile Submitted Successfully!")

        # BMI Calculation
        bmi = calculate_bmi(weight, height)
        bmi_status = get_category(bmi)

        st.write(f"### 📊 Your BMI: {bmi} ({bmi_status})")

        # Generate Plan
        with st.spinner("Generating your 5-day workout plan..."):
            prompt = build_prompt(
                name=name,
                height=height,
                weight=weight,
                goal=goal,
                level=level,
                equipment=equipment,
                bmi=bmi,
                bmi_status=bmi_status
            )

            full_plan = query_model(prompt)

        # Display Plan
        st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
        st.write(full_plan)