import streamlit as st
import os
import requests


@st.cache_data
def query(text):
    API_URL = "https://api-inference.huggingface.co/models/j-hartmann/emotion-english-distilroberta-base"
    headers = {"Authorization": "Bearer " + os.environ["HF_API_KEY"]}
    payload = {"inputs": text}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()


# Initialize the `prompts` session state
if "prompts" not in st.session_state:
    st.session_state.prompts = []

col1, col2 = st.columns(2)

with col1:
    if st.button("Add query"):
        st.session_state.prompts.append("")
    if st.button("Remove query"):
        st.session_state.prompts.pop(0)


def card(index: int):
    def on_change():
        # Update the `prompts[index]` session state with the value from the text input which is specified by the key, `f"prompt_{index}"`.
        st.session_state.prompts[index] = st.session_state[f"prompt_{index}"]
    st.text_input("Prompt:", key=f"prompt_{index}", value=st.session_state.prompts[index], on_change=on_change)

    prompt = st.session_state.prompts[index]

    # This `query` function is cached, so it will only be re-run if the input `prompt` changes.
    result = query(prompt)

    st.write(f'Prompt: "{prompt}"')
    st.json(result)


with col2:
    for index in range(len(st.session_state.prompts)):
        card(index)