import openai # OpenAI API
import os # To get the API key from the environment variables
import re # To check if input arguments are in the correct format
import ast # To convert string to array
import gradio as gr # To create the interface

openai.api_key = os.environ["OPENAI_API_KEY"]

#A function to generate scenes of a movie, given a description of the movie, and the characters of the movie.
def generate_scenes(input_description, character_names , character_descriptions):
    # Check if input arguments are in the correct format
    if (re.match(r'^\[\s*".*"\s*(,\s*".*"\s*)*\]$', character_names) and
            re.match(r'^\[\s*".*"\s*(,\s*".*"\s*)*\]$', character_descriptions)):
        raise ValueError("Invalid input format")

    # Convert character_names and character_descriptions to arrays
    character_names = ast.literal_eval(character_names)
    character_descriptions = ast.literal_eval(character_descriptions)

    characters_string = ""

    for i in range(len(character_names)):
        characters_string += character_descriptions[i]+"\n"

    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a clever movie script writer. Write the scenes for an original, romantic, and funny movie using the Save the Cat structure. Use the same language as the user after 'Place','Beat', and 'Plot element'."},
            {"role": "user", "content": "In a world of magic and myth, the ancient Dark Lord Sauron seeks to regain his lost power by reclaiming the One Ring, a powerful artifact that can enslave all of Middle-earth. The Ring, discovered by a young hobbit named Frodo Baggins, must be destroyed to prevent Sauron's rise to power once more.\nFrodo is the hero. A naive young hobbit, living a normal peaceful life until now.\nGandalf is the mentor figure. A white mage, he will lead the fellowship towards the destruction of the ring in the dark lands.\nSauron is the antagonist. He is the commander of the dark forces, extremely powerful, and capable of unimaginable things to control Middle Earth.\nSam, Frodo's best friend. He is a feisty and brave hobbit. He is a very loyal friend to Frodo, and will stay by his side no matter what.\nBoromir is a man from the fellowship who swore to protect Frodo but is slowly being seduced by the power of the ring, and ends up trying to steal it from him."},
            {"role": "assistant", "content": "<scenes>\nPlace: The shire.\nPlot element: Opening Image.\nBeat: Frodo is reading a book in the woods by himself.\nPlace: The Shire.\nPlot element: Theme Stated.\nBeat: Some believe that it is only great power that can hold evil in check. But that is not what Gandalf has found. He founds it is the small things. Everyday deeds by ordinary folks that keeps the darkness at bay.\nPlace: The Shire.\nPlot element: Set-up.\nBeat: Frodo is living a peaceful life with his uncle in The Shire.\nPlace: The Shire.\nPlot element: Catalyst.\nBeat: Bilbo Baggins wears the ring, and Gandalf finds out about it. Now the forces of darkness are after the ring and whoever holds it.\nPlace: The Shire.\nPlot element: Debate.\nBeat: Frodo argues with Gandalf about who should carry the ring. He argues he shouldn't carry it.\nPlace: The Shire.\nPlot element: Break Into Two.\nBeat: Gandalf convinces Frodo to take the ring into the land of darkness to destroy it.\nPlace: Middle Earth.\nPlot element: B Story is introduced.\nBeat: Frodo and his best friend Sam start the journey towards the land of darkness by themselves. Wary of all the dangers that they will find along the way. They. met the fellowship of the ring that consists of men, dwarves, and elves that will fight for the Middle Earth, and to protect Frodo.\nPlace: Middle Earth.\nPlot element: Fun & Games.\nBeat: The team is training very hard to fight with the malevolous orcs. The fellowship travels throughout the Middle Earth towards the dark lands.\nPlace: Middle Earth.\nPlot element: Midpoint.\nBeat: The fellowship faces a powerful dark resistance which ends up killing Gandalf. The team is devastated, but they must keep on going. They find out the journey will not be as easy as they were thinking.\nPlace: Isengard.\nPlot element: Bad Guys Close In.\nBeat: The dark lord Sauron creates a more powerful kind of orc to defeat the fellowship, and hunt down Frodo. The new orcs, called Uruk hai, begin their hunt for the ring. One of the men of the fellowship, named Boromir, is slowly being seduced by the power of the ring, and dreams of stealing it from Frodo.\nPlace: Amon Hen.\nPlot element: All Is Lost.\nBeat:Boromir tries to steal the ring from Frodo, but he is not successful in doing so. The fellowship is broken, and now Frodo is left to travel the rest of the journey alone. The orcs find them. \nPlace: Amon Hen.\nPlot element: Dark Night of the Soul.\nBeat: Frodo is afraid that someone else in the fellowship might steal the ring. He is sad, alone, and fearful. \nPlace: Amon Hen.\nPlot element: Break into Three.\nBeat: Frodo escapes from the orcs, and his fellowship. The rest of the fellowship decides to fight back the orcs.\nPlace: Amon Hen.\nPlot element: Finale\nBeat:Frodo manages to escape, followed by his best friend Sam. The fellowship defeats the orcs in an all out fight. Boromir fights to death with an orc, he loses.\nPlace: Amon Hen.\nPlot element: Final Image\nBeat: Frodo and his best friend escape and continue the journey on their own. The rest of the fellowship begin a revenge hunt for the last orcs after burying Boromir.<end>"},
            {"role": "user", "content": "En un pequeño pueblo, dos jóvenes de familias rivales se enamoraron inesperadamente al encontrarse en una fiesta. A pesar de las dificultades y secretismo, su amor creció fuerte y desafiante. Con el tiempo, lograron unir a sus familias y su amor se convirtió en leyenda.\nMaría es la amada de Juan. Es una joven dulce y amable. Ella pertenece a la familia rival de Juan.\nJuan es la pareja perfecta, es atento, caballeroso y servicial. Está perdidamente enamorado de María.\nLos padres de Juan son líderes de la familia rival. Ellos no están de acuerdo con la relación de Juan y María.\nLos padres de María son líderes de la familia rival. Ellos no están de acuerdo con la relación de María y Juan.\nLos amigos de Juan y María son un grupo de jóvenes que los apoyan en su amor. Ellos ayudan a Juan y María a mantener su relación en secreto."},
            {"role": "assistant", "content": "<scenes>\nPlace: Plaza principal de un pintoresco pueblo mexicano.\nPlot element: Imagen de apertura.\nBeat: Juan y María se encuentran en la fiesta y cruzan miradas, sintiendo una conexión instantánea.\nPlace: Calle empedrada junto a una colorida cantina.\nPlot element: Tema establecido.\nBeat: Los amigos de Juan y María hablan sobre cómo el amor verdadero puede superar cualquier obstáculo, incluso la rivalidad entre sus familias.\nPlace: Jardín secreto detrás de una antigua iglesia.\nPlot element: Preparación.\nBeat: Juan y María se enamoran y comienzan a encontrarse a escondidas, a pesar del descontento de sus familias.\nPlace: Mirador con vista al pueblo y al atardecer.\nPlot element: Catalizador.\nBeat: Juan y María se besan por primera vez, sellando su amor y su deseo de estar juntos a pesar de las dificultades.\nPlace: Callejón del pueblo.\nPlot element: Debate.\nBeat: Los amigos de Juan y María debaten si deben ayudar a la pareja a mantener su relación en secreto o si deben enfrentarse a sus familias y revelar la verdad.\nPlace: Casa de los amigos de Juan y María.\nPlot element: Entrada en dos.\nBeat: Los amigos deciden ayudar a Juan y María a mantener su relación en secreto mientras trabajan en un plan para unir a las dos familias rivales.\nPlace: Pequeño pueblo.\nPlot element: Subtrama introducida (Historia B).\nBeat: Los amigos de Juan y María comienzan a trabajar juntos para crear situaciones que muestren a las familias que tienen más en común de lo que piensan.\nPlace: Pequeño pueblo.\nPlot element: Diversión y juegos.\nBeat: Juan y María comparten momentos románticos y divertidos mientras sus amigos continúan trabajando para unir a las familias.\nPlace: Pequeño pueblo.\nPlot element: Punto medio.\nBeat: Las familias de Juan y María comienzan a sospechar que algo está sucediendo entre ellos, lo que provoca tensión y desconfianza.\nPlace: Pequeño pueblo.\nPlot element: Los malos se acercan.\nBeat: Los padres de Juan y María descubren la verdad sobre su relación y exigen que pongan fin a su amor.\nBeat: Pequeño pueblo.\nPlot element: Todo está perdido.\nBeat: Juan y María se ven obligados a separarse, dejando a ambos con el corazón roto y a sus amigos desanimados.\nPlace: Pequeño pueblo.\nPlot element: Noche oscura del alma.\nBeat: Juan y María reflexionan sobre su amor y lo que están dispuestos a hacer para estar juntos.\nPlace: Pequeño pueblo.\nPlot element: Entrada en tres.\nBeat: Juan y María deciden enfrentarse a sus familias y luchar por su amor, con la ayuda de sus amigos.\nPlace: Pequeño pueblo.\nPlot element: Final.\nBeat: Las familias de Juan y María, conmovidas por la fuerza de su amor, deciden dejar de lado sus diferencias y aceptar la relación.\nPlace: Pequeño pueblo.\nPlot element: Imagen final.\nBeat: Juan y María se casan, uniendo a sus familias y convirtiendo su amor en una leyenda en el pueblo.<end>"},
            {"role": "user", "content": input_description + "\n"+characters_string},
        ],
        temperature=0.5
    )
    scenes_places = []
    scenes_plot_elements = []
    scenes_beats = []
    output = response.choices[0].message.content
    for i in range(len(str(output).split("\n"))):
        if output.split("\n")[i].split(":")[0] == "Place":
            scenes_places.append(
                output.split("\n")[i].split(":")[1])
        elif output.split("\n")[i].split(":")[0] == "Plot element":
            scenes_plot_elements.append(
                output.split("\n")[i].split(":")[1])
        elif output.split("\n")[i].split(":")[0] == "Beat":
            scenes_beats.append(
                output.split("\n")[i].split(":")[1])
    
    #Write a string in markdown format with the scenes
    scenes_string = ""
    for i in range(len(scenes_places)):
        scenes_string += "## Beat " + str(i+1) + "\n"
        scenes_string += "**Place:** " + scenes_places[i] + "\n"
        scenes_string += "**Plot element:** " + scenes_plot_elements[i] + "\n"
        scenes_string += "**Beat:** " + scenes_beats[i] + "\n"
        scenes_string += "\n"

    return scenes_places, scenes_plot_elements, scenes_beats, scenes_string


demo = gr.Interface(
    fn=generate_scenes,
    # input
    inputs=[gr.Textbox(label="Descripción de la historia"),
            gr.Textbox(label="Nombres de los Personajes"),
            gr.Textbox(label="Descripción de los Personajes")],
    # output
    outputs=[
        # title
        gr.Textbox(label="Scenes Places"),
        gr.Textbox(label="Scenes Plot Elements"),
        gr.Textbox(label="Scenes Beats"),
        gr.Markdown(label="Scenes")],
    title="Dramatron Scenes 2.0",
    description="Escribe la descripción de la historia para la cuál quieres generar las escenas.")

demo.launch()