import streamlit as st
from utils import im_2_b64, RANDOM_SEED, client


with_prompt = """
In the following task, you will be presented with some images and a story that is, in some manner, related to that goal. You will be given a specific  goal and entity (generally, person or object), and be asked to identify condition necessary for that goal and the alternate condition that could prevent that goal.

Conditions for the output:
1. Condition: The condition is the necessary condition for the goal to be achieved. If the condition is not met, the goal cannot be achieved. The condition must be related to the entity.

2. Alternate Condition: The alternate condition is a different version of the condition that would prevent the goal from being achieved. It is likely that this alternate condition will contradict information provided in the images and/or story.

Output in a python dictionary where it should have the following keys: 'condition', 'alternate_condition'.

Story: {story}

Entity: {entity}

Goal: {goal}
"""

wo_prompt = """
In the following task, you will be presented with a image and a story that is, in some manner, related to that goal. You will be given a specific  goal, and be asked to identify an entity (person or object), condition necessary for that goal and the alternate condition that could prevent that goal.

Conditions for the output:
1. Entity: The entity is the person or object that the goal is related to. The entity should be a crucial part for achieving the goal.

2. Condition: The condition is the necessary condition for the goal to be achieved. If the condition is not met, the goal cannot be achieved. The condition must be related to the entity.

3. Alternate Condition: The alternate condition is a different version of the condition that would prevent the goal from being achieved. It is likely that this alternate condition will contradict information provided in the images and/or story.

Output in a python dictionary where it should have the following keys: 'entity', 'condition', 'alternate_condition'.

Story: {story}

Goal: {goal}
"""

def get_gpt4V_response_1(story, goal, entity, images, temperature=0.5):
    # Convert image to base64
    image_urls = []
    for image in images:
        if type(images[0]) == bytes:
            image_b64 = image
            image_url = f"data:image/jpeg;base64,{image_b64.decode('utf-8')}"
        else:
            image_b64 = im_2_b64(image)
            image_url = f"data:image/jpeg;base64,{image_b64.decode('utf-8')}"
        image_urls.append(image_url)
    st.write("✅ Image converted")

    if entity:
        prompt = with_prompt
        now_prompt = prompt.format(story=story, goal=goal, entity=entity)
    else:
        prompt = wo_prompt
        now_prompt = prompt.format(story=story, goal=goal)

    content = [
        {"type": "text", "text": now_prompt},
    ]

    st.write("✅ Prompt created")

    for image_url in image_urls:
        content.append({
            "type": "image_url",
            "image_url": {
                "url": image_url,
            },
        })
    
    st.write("🚀 Getting Response from GPT4V")
    response = client.chat.completions.create(
        model="gpt-4-vision-preview",
        seed=RANDOM_SEED,
        messages=[
            {
                "role": "user",
                "content": content
            }
        ],
        temperature=temperature,
        max_tokens=1024,
        # top_p=1,
        # frequency_penalty=0,
        # presence_penalty=0,
    )
    # print(response)
    # print("Prompt:")
    # print(now_prompt)
    out = response.choices[0].message.content
    # print("OUTPUT:", out)
    # print("====================================")
    # print()

    st.write("✅ Response generated")

    return out