import requests
import re
import base64
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from PIL import Image
from flask import Flask, render_template, request, redirect, url_for, flash

app = Flask(__name__)

# Load the Hugging Face model and tokenizer
model_id = "meta-llama/llama-3-2-90b-vision-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

def input_image_setup(uploaded_file):
    """
    Encodes the uploaded image file into a base64 string to be used with AI models.

    Parameters:
    - uploaded_file: File-like object uploaded via a file uploader

    Returns:
    - encoded_image (str): Base64 encoded string of the image data
    """
    if uploaded_file is not None:
        bytes_data = uploaded_file.read()
        encoded_image = base64.b64encode(bytes_data).decode("utf-8")
        return encoded_image
    else:
        raise FileNotFoundError("No file uploaded")

def format_response(response_text):
    """
    Formats the model response to display each item on a new line as a list.
    Converts numbered items into HTML `<ul>` and `<li>` format.
    """
    response_text = re.sub(r"\*\*(.*?)\*\*", r"<p><strong>\1</strong></p>", response_text)
    response_text = re.sub(r"(?m)^\s*\*\s(.*)", r"<li>\1</li>", response_text)
    response_text = re.sub(r"(<li>.*?</li>)+", lambda match: f"<ul>{match.group(0)}</ul>", response_text, flags=re.DOTALL)
    response_text = re.sub(r"</p>(?=<p>)", r"</p><br>", response_text)
    response_text = re.sub(r"(\n|\\n)+", r"<br>", response_text)
    return response_text

def generate_model_response(encoded_image, user_query, assistant_prompt):
    """
    Sends an image and a query to the model and retrieves the description or answer.
    Formats the response using HTML elements for better presentation.
    """
    # Prepare input for the model
    input_text = assistant_prompt + "\n\n" + user_query + "\n![Image](data:image/jpeg;base64," + encoded_image + ")"
    
    inputs = tokenizer(input_text, return_tensors="pt")
    
    try:
        # Generate the model's response
        outputs = model.generate(**inputs)
        raw_response = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Format the raw response text using the format_response function
        formatted_response = format_response(raw_response)
        return formatted_response
    except Exception as e:
        print(f"Error in generating response: {e}")
        return "<p>An error occurred while generating the response.</p>"

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        user_query = request.form.get("user_query")
        uploaded_file = request.files.get("file")

        if uploaded_file:
            encoded_image = input_image_setup(uploaded_file)

            if not encoded_image:
                flash("Error processing the image. Please try again.", "danger")
                return redirect(url_for("index"))

            assistant_prompt = """
            You are an expert nutritionist. Your task is to analyze the food items displayed in the image and provide a detailed nutritional assessment using the following format:

            1. **Identification**: List each identified food item clearly, one per line.
            2. **Portion Size & Calorie Estimation**: For each identified food item, specify the portion size and provide an estimated number of calories. Use bullet points with the following structure:
            - **[Food Item]**: [Portion Size], [Number of Calories] calories

            Example:
            *   **Salmon**: 6 ounces, 210 calories
            *   **Asparagus**: 3 spears, 25 calories

            3. **Total Calories**: Provide the total number of calories for all food items.

            Example:
            Total Calories: [Number of Calories]

            4. **Nutrient Breakdown**: Include a breakdown of key nutrients such as **Protein**, **Carbohydrates**, **Fats**, **Vitamins**, and **Minerals**. Use bullet points, and for each nutrient provide details about the contribution of each food item.

            Example:
            *   **Protein**: Salmon (35g), Asparagus (3g), Tomatoes (1g) = [Total Protein]

            5. **Health Evaluation**: Evaluate the healthiness of the meal in one paragraph.

            6. **Disclaimer**: Include the following exact text as a disclaimer:

            The nutritional information and calorie estimates provided are approximate and are based on general food data. 
            Actual values may vary depending on factors such as portion size, specific ingredients, preparation methods, and individual variations. 
            For precise dietary advice or medical guidance, consult a qualified nutritionist or healthcare provider.

            Format your response exactly like the template above to ensure consistency.
            """

            # Generate the model's response
            response = generate_model_response(encoded_image, user_query, assistant_prompt)

            # Render the result
            return render_template("index.html", user_query=user_query, response=response)

        else:
            flash("Please upload an image file.", "danger")
            return redirect(url_for("index"))

    return render_template("index.html")

if __name__ == "__main__":
    app.run(debug=True)