import gradio as gr
import numpy as np
from PIL import Image
import base64
from io import BytesIO
import json

from huggingface_hub import InferenceClient
from transformers import AutoProcessor

idefics_processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")


def process_images_and_text(image_path, query, client):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": query},
            ],
        },
    ]

    prompt_with_template = idefics_processor.apply_chat_template(
        messages, add_generation_prompt=True
    )

    # encode images to strings which can be sent to the endpoint
    def encode_local_image(image):
        pil_image = Image.fromarray(image.astype("uint8"), "RGB")

        # Convert the image to a base64 string
        buffer = BytesIO()
        pil_image.save(
            buffer, format="JPEG"
        )  # Use the appropriate format (e.g., JPEG, PNG)
        base64_image = base64.b64encode(buffer.getvalue()).decode("utf-8")

        # add string formatting required by the endpoint
        image_string = f"data:image/jpeg;base64,{base64_image}"

        return image_string

    image_string = encode_local_image(image_path)
    prompt_with_images = prompt_with_template.replace("<image>", "![]({}) ").format(
        image_string
    )

    payload = {
        "inputs": prompt_with_images,
        "parameters": {
            "return_full_text": False,
            "max_new_tokens": 200,
        },
    }

    return json.loads(client.post(json=payload).decode())[0]


# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def resize_image(img):
    width, height = img.size
    img = img.resize((int(width / 2), int(height / 2)))
    return img


idefics_client = InferenceClient("HuggingFaceM4/idefics2-8b-chatty")
rating_client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")


def ask_question_on_image(question: str, image: str) -> str:
    try:
        return process_images_and_text(image, question, idefics_client)
    except Exception as e:
        print(e)
        if "Payload Too Large" in str(e):
            image = resize_image(image)
            return process_images_and_text(image, question, idefics_client)


def give_rating(final_grade, description):
    prompt = f"You have been given a picture of a person to grade on a scale 0-10. The final grade is {final_grade}"
    if final_grade < 5:
        prompt += """Based on their appearance, the person is a terrible human being.
The user should absolutely understand that they should never ever be seen with that person.
It's not only that the person described does not 'look' cool, but they are a horrible person, a nuisance to society: they are ugly on the inside even more than on the outside.
Of course this is a humoristic message, and the criteria you use should be nonsensical funny, but the message below each criterion should appear serious, with a stern voice. Include words like "this person is a nuisance", "scum", "scourge", "tepid", "horrid", "rot on the surface of Earth", "never under any circumstance should you be seen with that person or you will instantly be drop down to the absolut zero level of coolness".
Make it excessively funny by telling funny takes like "the eyebrow-to-eye colour ratio is horrendous" with a very serious tone!"""
    else:
        prompt += """"Based on their appearance, the person has been rated as an outstanding individual! Describe the person as if you'd just fallen in love with them! You cannot overstate enough how much of a pure miracle that person is. Include words like "butterflies in my stomach", "impeccable haircut", "radiant sunshine", "Do you know this feeling of having a bad day ? This person certainly does not", etc."""
    prompt += f"Here is the person's description: {description}"
    prompt += "Now you have to justify the grade given above: please create a list of completely bullshit and funny criteria and justify why the grade is high or low on each criterion. Make sure to keep a serious tone while giving these bullshit reasons, the reader must not know that you're joking! Give at most 4 criteria, and only a few sentences in each: be concise!"
    messages = [
        {"role": "user", "content": prompt},
    ]
    output = (
        rating_client.chat_completion(messages=messages, max_tokens=1500)
        .choices[0]
        .message.content
    )
    return output


def rate_coolness(image1, image2):
    probabilities = [0.1, 0.1, 0.1, 0.2, 0.3, 0.2]
    numbers = [0, 1, 2, 9, 9.5, 10]
    grade = np.random.choice(numbers, p=probabilities)
    final_rating = f"### Final grade: **{grade}/10**\n"
    image_description = ask_question_on_image(
        f"Please describe this image in great detail: who does the person seems to be? What's her character? Is there anything noticeable in their appearance or clothing compared to a normal human being?",
        image1,
    )["generated_text"]

    final_rating += give_rating(grade, image_description)
    return final_rating


with gr.Blocks(theme=gr.themes.Soft(primary_hue="red")) as demo:
    image_1 = gr.Image()
    button = gr.Button("How cool is this person?")
    out = gr.Markdown()
    button.click(rate_coolness, inputs=[image_1], outputs=[out])


if __name__ == "__main__":
    demo.launch()