from PIL import Image
from io import BytesIO
import base64
import streamlit as st


UNIT_COST = {
    "user": 0.01,
    "assistant": 0.03,
}


# Convert Image to Base64 
def im_2_b64(image):
    image = Image.open(image)
    image.thumbnail((512, 512), Image.Resampling.LANCZOS)
    image = image.convert("RGB")
    buff = BytesIO()
    image.save(buff, format="JPEG")
    img_str = base64.b64encode(buff.getvalue())
    return img_str


def calculate_cost():
    def get_text_cost(text, unit_cost):
        num_of_words = len(text.split(" "))
        tokens = max(1000.0 * num_of_words / 750.0, 0.0)
        tokens = tokens / 1000.0
        cost = tokens * unit_cost
        return cost
    def get_image_cost(unit_cost=0.01):
        cost = 0.00255 # 512x512 image: https://openai.com/pricing
        return cost

    messages = st.session_state.messages
    total_cost = 0
    for message in messages:
        role = message["role"]
        for content in message["content"]:
            if content["type"] == "image_url":
                total_cost += get_image_cost(UNIT_COST[role])
            else:
                total_cost += get_text_cost(content["text"], UNIT_COST[role])
    
    st.session_state.cost.append(total_cost)