import streamlit as st
import google.generativeai as genai
from PIL import Image
import requests

# Configure Google Generative AI library with an API key
genai.configure(api_key="AIzaSyBuQHthxfR0t-3RtZC44YxTBJWN4roz-O4")

# API configurations for Dobby
DOBBY_API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
DOBBY_API_KEY = "fw_3ZjtsywUGddwa1wGY4VvB3eW"
DOBBY_LEASHED_MODEL = "accounts/sentientfoundation/models/dobby-mini-leashed-llama-3-1-8b#accounts/sentientfoundation/deployments/22e7b3fd"
DOBBY_UNHINGED_MODEL = "accounts/sentientfoundation/models/dobby-mini-unhinged-llama-3-1-8b#accounts/sentientfoundation/deployments/81e155fc"

# Initialize session state for tracking meals and user info
if "meal_data" not in st.session_state:
    st.session_state.meal_data = {"Breakfast": {}, "Lunch": {}, "Dinner": {}}
if "daily_totals" not in st.session_state:
    st.session_state.daily_totals = {"Calories": 0, "Protein": 0, "Carbs": 0, "Fat": 0}
if "user_info" not in st.session_state:
    st.session_state.user_info = {}


st.markdown("""
    <style>
        .stApp {
            background-color: #1a1a1a; /* Dark black background */
            color: #FFFFFF;  /* Set text color to white for readability */
            font-family: 'Inter', sans-serif; /* Inter font */
        }
        .title {
            font-size: 36px;
            font-weight: bold;
            color: #FFFFFF; /* White text for title */
            text-align: center;
            margin-bottom: 20px;
        }
        .subtitle {
            font-size: 20px;
            font-weight: 400;
            color: #CCCCCC; /* Light grey for subtitle */
            text-align: center;
            margin-bottom: 40px;
        }
        .header {
            font-size: 24px;
            font-weight: bold;
            color: #FFFFFF; /* White text for headers */
            margin-bottom: 10px;
        }
        .button {
            background-color: #4CAF50; /* Green button */
            color: white;
            font-size: 16px;
            padding: 10px 20px;
            border-radius: 5px;
        }
    </style>
""", unsafe_allow_html=True)



st.markdown('<div class="title">Calorie Bro</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Your personal AI image calorie tracker, powered by Sentient Dobby AI. Upload you breakfast, lunch and dinner, write your workout for the day and dare to get some 18+ feedback!</div>', unsafe_allow_html=True)

# Step 1: User inputs weight (in pounds), goal, and Dobby tone
if not st.session_state.user_info:
    st.markdown('<div class="header">Step 1: Set Your Preferences</div>', unsafe_allow_html=True)
    weight = st.number_input(
        "Enter your weight (lbs):", 
        min_value=50,  # Minimum weight in pounds (~30 kg)
        max_value=1000,  # Maximum weight in pounds (~300 kg)
        step=1, 
        value=175
    )
    goal = st.selectbox("Select your goal:", ["Lose Weight", "Maintain Weight", "Gain Weight"])
    tone = "Sarcastic (Unhinged)"  # Default tone
    if st.button("Submit", help="Save your preferences to start tracking meals."):
        st.session_state.user_info = {
            "Weight": weight,
            "Goal": goal,
            "Tone": tone
        }
        st.success("Your preferences have been saved!")

# Step 2: Meal tracking
if st.session_state.user_info:
    meal = st.selectbox("Select Meal:", ["Breakfast", "Lunch", "Dinner"])
    uploaded_file = st.file_uploader(f"Upload your {meal} image", type=["jpg", "jpeg", "png"])

    # Functions for analysis and comments
    def analyze_food_with_gemini(input_prompt, image_data):
        model = genai.GenerativeModel("gemini-1.5-pro-latest")
        response = model.generate_content([input_prompt, image_data[0]])
        return response.text

    def prepare_image_data(uploaded_file):
        bytes_data = uploaded_file.getvalue()
        return [{"mime_type": uploaded_file.type, "data": bytes_data}]

    def get_dobby_comment(food_name, total_calories, total_protein, total_carbs, total_fat, meal, goal, tone):
        selected_model = DOBBY_UNHINGED_MODEL  # Default tone is sarcastic
        dobby_prompt = f"""
        The user had {meal}: {food_name}.
        Nutritional data: {total_calories} kcal, {total_protein} g protein, {total_carbs} g carbs, {total_fat} g fat.
        The user's goal is to {goal.lower()}. Provide feedback for this meal.
        """
        response = requests.post(
            DOBBY_API_URL,
            headers={
                "Authorization": f"Bearer {DOBBY_API_KEY}",
                "Content-Type": "application/json",
            },
            json={
                "model": selected_model,
                "messages": [{"role": "user", "content": dobby_prompt}],
            },
        )
        return response.json().get("choices", [{}])[0].get("message", {}).get("content", "No response")

    # Process uploaded meal image
    if uploaded_file:
        st.image(uploaded_file, caption=f"Uploaded {meal} Image", use_container_width=True)
        with st.spinner("Analyzing the image..."):
            try:
                # Prepare image data and analyze with Gemini
                image_data = prepare_image_data(uploaded_file)
                input_prompt = """
                You are an expert nutritionist analyzing the food items in the image.
                Start by naming the meal based on the image, identify all ingredients, and estimate calories for each ingredient.
                Summarize the total calories, protein, carbs, and fat. Provide data in this format:
                Meal Name: [Name of the meal]
                Total Calories: X kcal
                Protein: X g
                Carbs: X g
                Fat: X g
                say nothing else
                """
                analysis = analyze_food_with_gemini(input_prompt, image_data)

                # Extract analysis data
                lines = analysis.split("\n")
                food_name = lines[0].replace("Meal Name: ", "").strip()
                total_calories = int(lines[1].replace("Total Calories: ", "").replace("kcal", "").strip())
                total_protein = int(lines[2].replace("Protein: ", "").replace("g", "").strip())
                total_carbs = int(lines[3].replace("Carbs: ", "").replace("g", "").strip())
                total_fat = int(lines[4].replace("Fat: ", "").replace("g", "").strip())

                # Save data for this meal
                st.session_state.meal_data[meal] = {
                    "Food Name": food_name,
                    "Calories": total_calories,
                    "Protein": total_protein,
                    "Carbs": total_carbs,
                    "Fat": total_fat,
                }

                # Update daily totals
                st.session_state.daily_totals["Calories"] += total_calories
                st.session_state.daily_totals["Protein"] += total_protein
                st.session_state.daily_totals["Carbs"] += total_carbs
                st.session_state.daily_totals["Fat"] += total_fat

                # Get Dobby's feedback
                dobby_comment = get_dobby_comment(
                    food_name, total_calories, total_protein, total_carbs, total_fat, meal,
                    st.session_state.user_info["Goal"], st.session_state.user_info["Tone"]
                )

                # Display results in the desired format
                st.write(
                    f"**{food_name}: {total_calories} kcal, {total_protein}g of Protein, {total_carbs}g of Carbs, {total_fat}g of Fat**"
                )
                st.write(dobby_comment)

            except Exception as e:
                st.error(f"Error processing the image: {e}")

# Step 3: Workout tracking and analysis
if all(st.session_state.meal_data[meal] for meal in ["Breakfast", "Lunch", "Dinner"]):
    st.subheader("Step 3: Log Your Workout")
    workout_input = st.text_area(
        "Describe your workout or activity today (e.g., 'I walked for 30 minutes and biked for 1 hour'):",
        placeholder="Enter your workout activities here..."
    )

    if st.button("Analyze Workout"):
        if workout_input.strip() == "":
            st.error("Please describe your workout before analyzing!")
        else:
            with st.spinner("Analyzing your workout..."):
                try:
                    # Gemini prompt for workout analysis
                    workout_prompt = f"""
                    You are an expert fitness tracker. The user did the following activities today:
                    {workout_input}
                    Calculate the total calories burned based on the description. 
                    Provide data in this format: "X calories burned by [activity descriptions]." Say nothing else.
                    """
                    # Analyze workout using Gemini for text
                    model = genai.GenerativeModel("gemini-1.5-pro-latest")
                    workout_analysis = model.generate_content([workout_prompt]).text.strip()
                    st.session_state.daily_totals["Workout"] = workout_analysis
    
                    # Dobby's feedback on the workout
                    dobby_workout_prompt = f"""
                    The user performed the following workout activities: {workout_input}.
                    Based on the analysis, {workout_analysis}.
                    Provide sarcastic or motivational feedback based on the user's goal to {st.session_state.user_info["Goal"].lower()}.
                    """
                    dobby_workout_comment = requests.post(
                        DOBBY_API_URL,
                        headers={
                            "Authorization": f"Bearer {DOBBY_API_KEY}",
                            "Content-Type": "application/json",
                        },
                        json={
                            "model": DOBBY_UNHINGED_MODEL,  # Default to unhinged tone
                            "messages": [{"role": "user", "content": dobby_workout_prompt}],
                        },
                    ).json().get("choices", [{}])[0].get("message", {}).get("content", "No response")
                    st.session_state.daily_totals["Workout Comment"] = dobby_workout_comment
    
                    # Display workout results
                    st.write(f"**{workout_analysis}**")
                    st.write(dobby_workout_comment)
                except Exception as e:
                    st.error(f"Error processing workout: {e}")


    # Step 4: Daily Summary
    if "Workout" in st.session_state.daily_totals:
        st.subheader("Daily Summary")
        st.write(f"**Total Calories Consumed:** {st.session_state.daily_totals['Calories']} kcal")
        st.write(f"**Total Protein:** {st.session_state.daily_totals['Protein']} g")
        st.write(f"**Total Carbs:** {st.session_state.daily_totals['Carbs']} g")
        st.write(f"**Total Fat:** {st.session_state.daily_totals['Fat']} g")
        st.write(f"**Workout Summary:** {st.session_state.daily_totals['Workout']}")

        # Final Dobby comment for the entire day
        final_prompt = f"""
        The user had the following meals today:
        Breakfast: {st.session_state.meal_data['Breakfast']['Food Name']}
        Lunch: {st.session_state.meal_data['Lunch']['Food Name']}
        Dinner: {st.session_state.meal_data['Dinner']['Food Name']}
        Total daily intake: {st.session_state.daily_totals['Calories']} kcal, 
        {st.session_state.daily_totals['Protein']} g protein, {st.session_state.daily_totals['Carbs']} g carbs, 
        {st.session_state.daily_totals['Fat']} g fat.
        They also did the following workout: {st.session_state.daily_totals['Workout']}
        Provide a summary of their day, considering both diet AND workout, and feedback for improvement.
        """
        final_comment = get_dobby_comment(
            "", st.session_state.daily_totals["Calories"], st.session_state.daily_totals["Protein"],
            st.session_state.daily_totals["Carbs"], st.session_state.daily_totals["Fat"], "Daily Summary",
            st.session_state.user_info["Goal"], st.session_state.user_info["Tone"]
        )
        st.write(final_comment)