import gradio as gr
import openai
import os as os

# OPENAI
openai.api_key = os.environ.get('openai-api')

#Functions
def generate_story(character1, character2, character3, character4, localizer, agefor, language, wordstouse, content):
    agefor = min(max(agefor, 2), 12)  # Ensure agefor is between 0 and 12
    
    # Define the prompt
    prompt = f"""
A story about {character1}, {character2}, {character3}, and {character4}, set in {localizer}.
The story is intended for a child aged {agefor}.' \
As much as possible use words that are in this list: {wordstouse}.
Make it fun and exciting with high burstiness and perplexity. Nobody wants a boring story :)
Integrate all the following elements in the story: {content}.
The story should be taught in {language}!
"""
    
    inject = [{"role":"user", "content":f"{prompt}"}]
    
    # Define the request to the OpenAI API
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=inject,
        temperature=0.8,
        max_tokens=1000,
        top_p=1,
        n=1,
        frequency_penalty=0,
        presence_penalty=0,
        stop=None
    )
    
    # Return the story
    return response['choices'][0]['message']['content']

with gr.Blocks() as demo:
    gr.Markdown("Enter the details below and then click **Make me a story** to generate your story.")
    
    with gr.Row():
        character1 = gr.Textbox(label="Character 1")
        character2 = gr.Textbox(label="Character 2")
        character3 = gr.Textbox(label="Character 3")
        character4 = gr.Textbox(label="Character 4")
    with gr.Row():
        content = gr.Textbox(label="What you want to hear about in this story", placeholder= "a giant elephant eating pink banana, a broken chair, ...")
    with gr.Row():
        wordstouse = gr.Textbox(label="Words for reading", info="Put here the words you would like the child to train on, not compulstory", placeholder= "the then car if this ...")
    with gr.Row():
        localizer = gr.Textbox(label="Location", value="Hong Kong", placeholder="Location")
        agefor = gr.Slider(minimum=2, maximum=12, step=1, value=6, label="Age of the child")
        language = gr.Dropdown(choices=["Mandarin", "Cantonese", "French", "English", "Korean", "Italian", "German", "Russian"], label="Language of the story", value="English")
    with gr.Row():
        btn = gr.Button("Make me a bed time story")
    with gr.Row():
        resultoutput = gr.Textbox(label="Here's my story")
        
    btn.click(fn=generate_story, inputs=[character1, character2, character3, character4, localizer, agefor, language, wordstouse, content], outputs=resultoutput)

demo.launch()