from typing import Tuple

import streamlit as st
import tiktoken


def num_tokens_from_messages(messages: list, model: str="gpt-3.5-turbo-0301") -> Tuple[int, int]:
    """Returns the number of tokens used by a list of messages."""
    try:
        encoding = tiktoken.encoding_for_model(model)
    except KeyError:
        st.sidebar.warning("Warning: model not found. Using cl100k_base encoding.")
        encoding = tiktoken.get_encoding("cl100k_base")
    if model == "gpt-3.5-turbo":
        st.sidebar.warning("Warning: gpt-3.5-turbo may change over time. Returning num tokens assuming gpt-3.5-turbo-0301.")
        return num_tokens_from_messages(messages, model="gpt-3.5-turbo-0301")
    elif model == "gpt-4":
        st.sidebar.warning("Warning: gpt-4 may change over time. Returning num tokens assuming gpt-4-0314.")
        return num_tokens_from_messages(messages, model="gpt-4-0314")
    elif model == "gpt-4-32k":
        st.sidebar.warning("Warning: gpt-4 may change over time. Returning num tokens assuming gpt-4-0314.")
        return num_tokens_from_messages(messages, model="gpt-4-0314")
    elif model == "gpt-3.5-turbo-0301":
        tokens_per_message = 5  # every message follows <|start|>{role/name}\n{content}<|end|>\n
    elif model == "gpt-4-0314":
        tokens_per_message = 4
    else:
        raise NotImplementedError(f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens.""")
    prompt_tokens = completion_tokens = 0
    for message in messages:
        if message['role'] in {'system', 'user'}:
            prompt_tokens += len(encoding.encode(message['content'])) + tokens_per_message
        else:
            completion_tokens += len(encoding.encode(message['content']))
    prompt_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
    return prompt_tokens, completion_tokens


def get_price(model: str) -> Tuple[float, float]:
    """$ / 1K tokens"""
    pricing = {
        'gpt-4': (0.03, 0.06),
        'gpt-4-32k': (0.06, 0.12),
        'gpt-3.5-turbo': (0.002, 0.002),
    }
    return pricing.get(model)


st.session_state["model"] = st.sidebar.selectbox(
    "Please select a model",
    ["gpt-3.5-turbo", "gpt-4", "gpt-4-32k"],
    help="ID of the model to use",
)
if 'messages' not in st.session_state:
    st.session_state['messages'] = [{"role": "system", "content": 'You are a helpful assistant'}]

st.title('ChatGPT Token Calculator')
metric_zone = st.empty()
input_zone = st.empty()

add_button = st.button('Add message', use_container_width=True, on_click = lambda: st.session_state['messages'].append({"role": "user", "content": ''}))


def change_role(index: int):
    st.session_state['messages'][index]['role'] = st.session_state['role' + str(index)]
    st.session_state['messages'][index]['content'] = st.session_state['content' + str(index)]


with input_zone.container():
    for index, message in enumerate(st.session_state['messages']):
        cols = st.columns([2, 5, 1])
        selections = ['user', 'system', 'assistant']
        role_index = selections.index(message['role'])
        # cols[0].selectbox('role', selections, key='role' + str(index), index=role_index, label_visibility='collapsed', on_change=lambda: st.session_state['messages'][index].__setitem__('role', st.session_state['role' + str(index)]))
        cols[0].selectbox('role', selections, key='role' + str(index), index=role_index, label_visibility='collapsed', on_change=change_role, args=(index, ))
        # cols[1].text_input('content', value=message.get('content'), key='content' + str(index), placeholder='Content',label_visibility='collapsed',on_change=lambda: st.session_state['messages'][index].__setitem__('content',st.session_state['content' + str(index)]))
        cols[1].text_input('content', value=message.get('content'), key='content' + str(index), placeholder='Content',label_visibility='collapsed', on_change=change_role, args=(index, ))
        cols[2].button('❌', key='remove'+str(index), on_click=lambda: st.session_state['messages'].pop(index))

with metric_zone.container():
    prompt_tokens, completion_tokens = num_tokens_from_messages(st.session_state['messages'], st.session_state['model'])
    prompt_price, completion_price = get_price(st.session_state['model'])
    col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
    col1.metric(label='Total Tokens', value=prompt_tokens+completion_tokens)
    col2.metric(label='Prompt Tokens', value=prompt_tokens)
    col3.metric(label='Completion Tokens', value=completion_tokens)
    col4.metric(label='Price', value=f'${(prompt_price*prompt_tokens+completion_price*completion_tokens)/1000:f}')