from openai import OpenAI
import streamlit as st
from utils import im_2_b64, calculate_cost, clear_uploader, undo, restart
import pickle
from upload import upload_file, get_file

share_keys = ["messages", "cost"]

client = OpenAI(api_key=st.secrets["OPENAI_KEY"])

st.set_page_config(
    page_title="ChatGPT with Vision",
    page_icon="🤖",
    menu_items={
        "About": """
            *ChatGPT with Vision* is a chat interface that uses OpenAI's API to generate responses to your prompts.  
            *---- Developed by **[Shubhashis Roy Dipta](https://roydipta.com)** ----*
        """
    }
)

if "messages" not in st.session_state:
    st.session_state.messages = []
if "uploader_key" not in st.session_state:
    st.session_state["uploader_key"] = 0
if "cost" not in st.session_state:
    st.session_state.cost = []

if len(st.session_state.messages) == 0 and "id" in st.query_params:
    with st.spinner("Loading chat..."):
        id = st.query_params["id"]
        data = get_file(id, 'chatgpt-vision-007')
        obj = pickle.loads(data)
        for k, v in obj.items():
            st.session_state[k] = v


def share():
    obj = {}
    for k in share_keys:
        if k in st.session_state:
            obj[k] = st.session_state[k]
    data = pickle.dumps(obj)
    id = upload_file(data, 'chatgpt-vision-007')
    url = f"https://umbc-nlp-chatgpt-vision.hf.space/?id={id}"
    st.success(f"Share URL: {url}")

with st.sidebar:
    st.title(":blue[ChatGPT with Vision]")

    password = st.text_input("Password", type="password")

    if st.button("Share", use_container_width=True):
        share()

    cols = st.columns(2)
    with cols[0]:
        if st.button("Restart", type="primary", use_container_width=True):
            restart()
    
    with cols[1]:
        if st.button("Undo", use_container_width=True):
            undo()

    with st.expander("Advanced Configuration"):
        st.subheader("Temperature")
        temperature = st.slider(label="x", min_value=0.1, max_value=1.0, value=0.2, step=0.1, label_visibility='collapsed')
        st.subheader("Max Tokens")
        max_tokens = st.slider(label="x", min_value=32, max_value=1024, value=256, step=32, label_visibility='collapsed')
        st.subheader("Random Seed")
        random_seed = st.number_input("Seed", min_value=0, max_value=1000000, value=42, step=1, label_visibility='collapsed')

    with st.expander("Image Input", expanded=True):
        images = st.file_uploader(
            "Image Upload",
            accept_multiple_files=True,
            type=["png", "jpg", "jpeg"],
            key=st.session_state["uploader_key"],
            label_visibility="collapsed",
        )

    with st.expander(f"Total Cost: ${sum(st.session_state.cost):.10f}"):
        if len(st.session_state.cost) > 0:
            st.subheader("Cost Breakdown")
            for i, c in enumerate(st.session_state.cost):
                st.write(f"Message {i+1}: ${c:.10f}")
            st.markdown("---")
            st.write(f"Total: ${sum(st.session_state.cost):.10f}")
        else:
            st.write("No cost incurred yet")

    append = st.checkbox("Append to previous message", value=False)


for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        contents = message["content"]
        
        for content in contents:
            if content["type"] == "text":
                st.markdown(content["text"])

        number_of_images = sum(1 for c in contents if c["type"] == "image_url")
        if number_of_images > 0:
            cols = st.columns(number_of_images)
            i = 0
            for content in contents:
                if content["type"] == "image_url":
                    with cols[i]:
                        st.image(content["image_url"]["url"])
                        i += 1


def push_message(role, content, images=None):
    contents = []
    contents.append({"type": "text", "text": content})
    if images:
        for image in images:
            image_b64 = im_2_b64(image)
            image_url = f"data:image/jpeg;base64,{image_b64.decode('utf-8')}"
            obj = {
                "type": "image_url",
                "image_url": {
                    "url": image_url,
                },
            }
            contents.append(obj)

    message = {"role": role, "content": contents}
    st.session_state.messages.append(message)
    return message

if prompt := st.chat_input("Type a message", key="chat_input"):
    if password != st.secrets["PASSWORD"]:
        st.error("Invalid password. Demo is read-only.")
        st.stop()
    push_message("user", prompt, images)
    with st.chat_message("user"):
        st.markdown(prompt)
        if images:
            cols = st.columns(len(images))
            for i, image in enumerate(images):
                with cols[i]:
                    st.image(image)

    if not append:
        with st.chat_message("assistant"):
            messages = [
                {"role": m["role"], "content": m["content"]}
                for m in st.session_state.messages
            ]
            stream = client.chat.completions.create(
                model="gpt-4-vision-preview",
                messages=messages,
                stream=True,
                seed=random_seed,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            response = st.write_stream(stream)

        push_message("assistant", response)
        calculate_cost()
    clear_uploader()