import streamlit as st
import base64
from huggingface_hub import InferenceClient
import os

# Initialize Hugging Face Inference client using token from environment variables
client = InferenceClient(api_key=os.getenv("HF_API_TOKEN"))

# 1. Function to identify dish from image
def identify_dish(image_bytes):
    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
    dish_name = ""

    try:
        for message in client.chat_completion(
            model="meta-llama/Llama-3.2-11B-Vision-Instruct",
            messages=[
                {
                    "role": "You are a food identification expert who identifies dishes from images. Your task is to strictly return the names of the dishes present in the image. Only return the dish names if you have high Confidence Level and without additional explanation or description.",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}" }},
                        {"type": "text", "text": "Identify the dishes in the image and return only the names of the dishes."},
                    ],
                }
            ],
            max_tokens=70,
            stream=True,
        ):
            if message.choices[0].delta.content:
                dish_name += message.choices[0].delta.content
    except:
        st.error("Error processing the image. Please upload a valid image.")
    return dish_name.strip()

# 2. Function to calculate metrics
def calculate_metrics(age, gender, height_cm, weight_kg, weight_goal, activity_level):
    bmi = weight_kg / ((height_cm / 100) ** 2)

    if gender == "male":
        bmr = 10 * weight_kg + 6.25 * height_cm - 5 * age + 5
    else:
        bmr = 10 * weight_kg + 6.25 * height_cm - 5 * age - 161

    activity_multipliers = {
        "sedentary": 1.2,
        "light": 1.375,
        "moderate": 1.55,
        "active": 1.725,
        "very active": 1.9
    }
    tdee = bmr * activity_multipliers[activity_level]

    if gender == "male":
        ibw = 50 + (0.91 * (height_cm - 152.4))
    else:
        ibw = 45.5 + (0.91 * (height_cm - 152.4))

    if weight_goal == "loss":
        daily_caloric_needs = tdee - 500
    elif weight_goal == "gain":
        daily_caloric_needs = tdee + 500
    else:
        daily_caloric_needs = tdee

    protein_calories = daily_caloric_needs * 0.2
    fat_calories = daily_caloric_needs * 0.25
    carbohydrate_calories = daily_caloric_needs * 0.55

    return {
        "BMI": bmi,
        "BMR": bmr,
        "TDEE": tdee,
        "IBW": ibw,
        "Daily Caloric Needs": daily_caloric_needs,
        "Protein Calories": protein_calories,
        "Fat Calories": fat_calories,
        "Carbohydrate Calories": carbohydrate_calories
    }

# 3. Function to generate diet plan
def generate_diet_plan(dish_name, calorie_intake_per_day, goal):
    user_input = f"""
    You are a certified Dietitian with 20 years of experience. Based on the following input, create an Indian diet plan that fits within the calculated calorie intake and assesses if the given dish is suitable for the user's goal.

    Input:
    - Dish Name: {dish_name}
    - Caloric Intake per Day: {calorie_intake_per_day} calories
    - Goal: {goal}
    """
    response = client.chat_completion(
        model="meta-llama/Meta-Llama-3-8B-Instruct",
        messages=[{"role": "You are a certified Dietitian with 20 years of Experience", "content": user_input}],
        max_tokens=500
    )

    return response.choices[0].message.content

# Streamlit Sidebar for user input
st.sidebar.title("User Input")
image_file = st.sidebar.file_uploader("Upload an image of the dish", type=["jpeg", "png"])
age = st.sidebar.number_input("Enter your age", min_value=1)
gender = st.sidebar.selectbox("Select your gender", ["male", "female"])
height_cm = st.sidebar.number_input("Enter your height (cm)", min_value=1.0)
weight_kg = st.sidebar.number_input("Enter your weight (kg)", min_value=1.0)
weight_goal = st.sidebar.selectbox("Weight goal", ["loss", "gain", "maintain"])
activity_level = st.sidebar.selectbox("Activity level", ["sedentary", "light", "moderate", "active", "very active"])
time_frame = st.sidebar.number_input("Time frame to achieve goal (months)", min_value=1)

# Submit Button
if st.sidebar.button("Submit"):
    if image_file:
        # Step 1: Identify the dish
        image_bytes = image_file.read()
        dish_name = identify_dish(image_bytes)
        if dish_name:
            st.write("### Results")
            st.divider()
            st.write("**Dish Name Identified:**")
            st.success(dish_name)

            # Step 2: Perform Calculations
            metrics = calculate_metrics(age, gender, height_cm, weight_kg, weight_goal, activity_level)
            st.divider()
            st.write("**Calculated Metrics:**")
            st.write(f"**BMI:** {metrics['BMI']:.2f}")
            st.write(f"**BMR:** {metrics['BMR']:.2f} calories")
            st.write(f"**TDEE:** {metrics['TDEE']:.2f} calories")
            st.write(f"**Ideal Body Weight:** {metrics['IBW']:.2f} kg")
            st.write(f"**Daily Caloric Needs:** {metrics['Daily Caloric Needs']:.2f} calories")

            # Step 3: Generate Diet Plan
            st.divider()
            st.write("**Diet Plan Based on Dish & Goal:**")
            diet_plan = generate_diet_plan(dish_name, metrics["Daily Caloric Needs"], weight_goal)
            st.info(diet_plan)
        else:
            st.error("Unable to identify the dish. Please try again with a different image.")
    else:
        st.error("Please upload an image of the dish.")

# CSS for styling
st.markdown("""
    <style>
    .stButton button { background-color: #4CAF50; color: white; }
    .stAlert { background-color: #f9f9f9; border: 1px solid #ddd; padding: 10px; }
    .stContainer { border: 1px solid #ddd; padding: 20px; margin-bottom: 20px; }
    .stSuccess { background-color: #e8f5e9; }
    .stInfo { background-color: #e3f2fd; }
    .stError { background-color: #ffebee; }
    </style>
""", unsafe_allow_html=True)