import json
import openai
import pandas as pd
import streamlit as st

st.set_page_config(layout="wide")

st.header('Calories')
openai.api_key = st.secrets["open_ai_key"]

if 'gpt_response' not in st.session_state:
    st.session_state.gpt_response = None

context = '''create a valid JSON array of objects for tracking the calorie and macronutrient content of specified food items in the following format:
[{
"food_item_name": "the name of the food item that was inputted including quantity expressed as a string",
"number_of_calories": "number of calories for the specified quantity of the food item expressed as an integer",
"grams_of_protein": "number of grams of protein for the specified quantity of the food item expressed as an integer",
"grams_of_fat": "number of grams of fat for the specified quantity of the food item expressed as an integer",
"grams_of_carbs": "number of grams of carbohydrates for the specified quantity of the food item expressed as an integer"
}]'''

prompt = st.text_input('List food items:')
st.write('Example: *1 cup broccoli, 500g boneless skinless chicken breast, 1 cup coffee*')

def get_response(context, prompt):
    st.session_state.gpt_response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": context},
            {"role": "user", "content": prompt}
        ],
        temperature=0.2,
        max_tokens=1000
    )
    
st.button(label='Submit', on_click=get_response, kwargs=dict(context=context, prompt=prompt))

if st.session_state.gpt_response is not None:
    st.dataframe(pd.DataFrame(json.loads(st.session_state.gpt_response['choices'][0]['message']['content'])), hide_index=True)
    cost = st.session_state.gpt_response['usage']["prompt_tokens"]*(0.0015/1000) + st.session_state.gpt_response['usage']["completion_tokens"]*(0.002/1000)
    st.write(f'Cost for query was approx ${cost}')