import gradio as gr
import openai
import os

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


def generate_characters(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. 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="Personajes")],
    title="Dramatron 2.0",
    description="Escribe la descripción de la historia para la cuál quieres generar los personajes.")

demo.launch()