import gradio as gr
import openai
import os

openai.api_key = "sk-"+os.environ['OPENAI_API_KEY']


def generate_characters(input_description):
    print(input_description)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "Write the characters for a given story using the same language as the user. You will be creative in order to find all the main characters, secondary, and tertiary characters so that the story becomes richer. The only part that should be in English is the separators between characters, which are <character>, <description>, <stop> and <end>."},
            {"role": "user", "content": "A science-fiction fantasy about a naive but ambitious farm boy from a backwater desert who discovers powers he never knew he had when he teams up with a feisty princess, a mercenary space pilot and an old wizard warrior to lead a ragtag rebellion against the sinister forces of the evil Galactic Empire."},
            {"role": "assistant", "content": "<character>Luke Skywalker <description>Luke Skywalker is the hero. A naive farm boy, he will discover special powers under the guidance of mentor Ben Kenobi.<stop>\n<character>Ben Kenobi <description>Ben Kenobi is the mentor figure. A recluse Jedi warrior, he will take Luke Skywalker as apprentice.<stop>\n<character>Darth Vader <description>Darth Vader is the antagonist. As a commander of the evil Galactic Empire, he controls space station The Death Star.<stop>\n<character>Princess Leia <description>Princess Leia is a feisty and brave leader of the Rebellion. She holds the plans of the Death Star. She will become Luke’s friend.<stop>\n<character>Han Solo <description>Han Solo is a brash mercenary space pilot of the Millenium Falcon and a friend of Chebacca. He will take Luke on his spaceship.<stop>\n<character>Chewbacca <description>Chewbacca is a furry and trustful monster. He is a friend of Han Solo and a copilot on the Millemium Falcon.<stop>\n<end>"},
            {"role": "user", "content": input_description}
        ],
        temperature=0.6
    )
    character_names = []
    character_descriptions = []
    content_split = response.choices[0].message.content.split("<stop>")
    for i in range(0, len(content_split) - 1):
        if content_split[i] == "<end>":
            break
        character_parts = content_split[i].split("<description>")
        if len(character_parts) >= 2:
            character_names.append(character_parts[0].strip())
            character_descriptions.append(character_parts[1].strip())
        else:
            character_names.append(character_parts[0].strip())
            # Add an empty description if not available
            character_descriptions.append("")
    return character_names, character_descriptions


demo = gr.Interface(
    fn=generate_characters,
    # input
    inputs=gr.Textbox(label="Descripción de la historia"),
    # output
    outputs=[
        # title
        gr.Textbox(label="Character Names"),
        gr.Textbox(label="Character Descriptions")],
    title="Dramatron Characters 2.0",
    description="Escribe la descripción de la historia para la cuál quieres generar los personajes.")

demo.launch()