from PIL import Image
from io import BytesIO
import base64
import time
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)


def clear_uploader():
    st.session_state["uploader_key"] += 1
    st.rerun()

def undo():
    if len(st.session_state.messages) > 0:
        st.query_params.clear()
        msg = st.session_state.messages.pop()
        if msg["role"] == "assistant":
            st.session_state.cost.pop()
        time.sleep(0.1)
        st.rerun()

def restart():
    st.query_params.clear()
    st.session_state.messages = []
    st.session_state.cost = []
    time.sleep(0.2)
    clear_uploader()