import streamlit as st
import requests

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

# ---------------------------------------------------
# HUGGING FACE API CONFIG
# ---------------------------------------------------

API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"

headers = {
    "Authorization": f"Bearer {st.secrets['HF_TOKEN']}",
    "Content-Type": "application/json"
}

def query(payload):
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)

        if response.status_code != 200:
            return {"error": f"API Error {response.status_code}: {response.text}"}

        return response.json()

    except requests.exceptions.RequestException as e:
        return {"error": f"Request failed: {str(e)}"}

# ---------------------------------------------------
# TITLE
# ---------------------------------------------------

st.title("🏋️ FitPlan AI – User Fitness Profile")

# ---------------------------------------------------
# PERSONAL INFORMATION
# ---------------------------------------------------

st.subheader("👤 Personal Information")

name = st.text_input("Enter Your Name")

gender = st.radio(
    "Gender",
    ["Male", "Female"],
    horizontal=True
)

# ---------------------------------------------------
# HEIGHT & WEIGHT
# ---------------------------------------------------

col1, col2 = st.columns(2)

with col1:
    height = st.number_input(
        "Height (in cm)",
        min_value=0.0,
        max_value=250.0,
        value=0.0,
        step=0.1
    )

with col2:
    weight = st.number_input(
        "Weight (in kg)",
        min_value=0.0,
        max_value=200.0,
        value=0.0,
        step=0.1
    )

# ---------------------------------------------------
# BMI FUNCTION
# ---------------------------------------------------

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

bmi = None

if height > 0 and weight > 0:
    height_m = height / 100
    bmi = weight / (height_m ** 2)

    st.metric("📊 Your BMI", f"{bmi:.2f}")
    st.info(f"BMI Category: {bmi_category(bmi)}")

# ---------------------------------------------------
# FITNESS GOAL
# ---------------------------------------------------

st.subheader("🎯 Fitness Goal")

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

# ---------------------------------------------------
# EQUIPMENT
# ---------------------------------------------------

st.subheader("🏋️ Available Equipment")

equipment_map = {}

col1, col2, col3 = st.columns(3)

with col1:
    equipment_map["No Equipment"] = st.checkbox("No Equipment")
    equipment_map["Pull-up Bar"] = st.checkbox("Pull-up Bar")
    equipment_map["Dip Bars"] = st.checkbox("Dip Bars")
    equipment_map["Push-up Handles"] = st.checkbox("Push-up Handles")
    equipment_map["Dumbbells"] = st.checkbox("Dumbbells")
    equipment_map["Adjustable Dumbbells"] = st.checkbox("Adjustable Dumbbells")

with col2:
    equipment_map["Barbell"] = st.checkbox("Barbell")
    equipment_map["Weight Plates"] = st.checkbox("Weight Plates")
    equipment_map["Kettlebells"] = st.checkbox("Kettlebells")
    equipment_map["Medicine Ball"] = st.checkbox("Medicine Ball")
    equipment_map["Yoga Mat"] = st.checkbox("Yoga Mat")
    equipment_map["Resistance Band"] = st.checkbox("Resistance Band")

with col3:
    equipment_map["Bosu Ball"] = st.checkbox("Bosu Ball")
    equipment_map["Stability Ball"] = st.checkbox("Stability Ball")
    equipment_map["Foam Roller"] = st.checkbox("Foam Roller")
    equipment_map["Treadmill"] = st.checkbox("Treadmill")
    equipment_map["Exercise Bike"] = st.checkbox("Exercise Bike")
    equipment_map["Skipping Rope"] = st.checkbox("Skipping Rope")

equipment = [item for item, selected in equipment_map.items() if selected]

# ---------------------------------------------------
# FITNESS LEVEL
# ---------------------------------------------------

st.subheader("📈 Fitness Level")

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

# ---------------------------------------------------
# SUBMIT BUTTON
# ---------------------------------------------------

if st.button("🚀 Submit Profile"):

    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_status = bmi_category(bmi)
        equipment_list = ", ".join(equipment)

        prompt = f"""<s>[INST]
You are a certified professional fitness trainer.

Generate a structured 5-day workout plan.

Rules:
- Clearly label Day 1 to Day 5
- Include exercises
- Include sets and reps
- Include short rest guidance
- Return only the workout plan

User Details:
Gender: {gender}
BMI: {bmi:.2f} ({bmi_status})
Goal: {goal}
Fitness Level: {fitness_level}
Equipment Available: {equipment_list}
[/INST]"""

        with st.spinner("Generating your AI workout plan..."):

            output = query({
                "inputs": prompt,
                "parameters": {
                    "max_new_tokens": 600,
                    "temperature": 0.3,
                    "return_full_text": False
                }
            })

            if isinstance(output, dict) and "error" in output:
                response = f"⚠️ {output['error']}\n\nPlease try again."

            elif isinstance(output, list):
                response = output[0]["generated_text"].strip()

                if "[/INST]" in response:
                    response = response.split("[/INST]")[-1].strip()

            else:
                response = "Unexpected response from model."

        st.subheader("🏋️ Your Personalized Workout Plan")
        st.markdown(response)