import logging
import os
import time

import cohere
import streamlit as st

DEFAULT_DISH = "herring and banana quesadillas"
MAX_TOKENS = 300

logger = logging.getLogger('recipe-logger')


def command(co, prompt, model='command-xlarge-20221108', max_tokens=MAX_TOKENS, *args, **kwargs):
    gens = co.generate(prompt, model=model, max_tokens=max_tokens, return_likelihoods='GENERATION', *args, **kwargs)
    return sorted(gens, key=lambda g: -g.likelihood)


def space(num_lines=1):
    for _ in range(num_lines):
        st.write("")


st.set_page_config(layout="centered", page_icon="🍔", page_title="Recipe Generation Using Cohere")

# Check if COHERE_API_KEY is not set from secrets.toml or os.environ
if "CO_API_KEY" not in os.environ:
    raise KeyError("CO_API_KEY not found in st.secrets or os.environ. Please set it in "
                   ".streamlit/secrets.toml or as an environment variable.")

# Adding a header to direct users to sign up for Cohere, explore the playground,
# and check out our git repo.
st.header("Recipe Generation")

#st.selectbox("Choose a cuisine:", options=['mexican', 'italian'], key="cuisine")

st.write("What would you like a recipe for?")
st.session_state.disabled = False

with st.form("form"):

    dish = st.text_input("Dish", placeholder=DEFAULT_DISH)
    options = st.multiselect("Special requests", [
        "extra delicious", "vegan", "gluten-free", "sugar-free", "dairy-free", "nut-free", "low-calorie", "low-fat",
        "low-sugar"
    ])
    submitted = st.form_submit_button(label="Generate Recipe!")

if submitted:
    with st.spinner("Cooking up something delicious ... please wait a few seconds ... "):
        dish = (dish or DEFAULT_DISH).strip()
        co = cohere.Client(api_key=os.environ['CO_API_KEY'])
        prompt = f"Give a recipe for {dish}."
        if options:
            prompt += f" Make sure it is {', '.join(options)}."
        logger.info(f"Sending prompt {repr(prompt)}")
        start = time.time()
        gens = command(co, prompt)
        recipe = gens[0].text
        num_tokens = len(gens[0].token_likelihoods)
        logger.info(f"Generated recipe with {num_tokens} tokens in {time.time()-start:.1f}s: {repr(recipe)}")
        st.markdown(f"## Recipe for {' '.join(options)} {dish}")
        st.write(recipe)
        st.markdown("## Delicious! 10/10 with rice")

# Draw chat history.
with st.expander("About", expanded=False):
    st.markdown(
        "This demo uses the Cohere API. Cohere provides access to advanced Large Language Models and NLP tools through one easy-to-use API. "
        "[**Get started for free!**]"
        "(https://dashboard.cohere.ai/welcome/register?utm_source=cohere-owned&utm_"
        "medium=content&utm_campaign=sandbox&utm_term=streamlit&utm_content=conversant)")