import gradio as gr
import openai
from reportlab.pdfgen import canvas

# FUNCTION


def storyGPT(key, name, situation, direction):
    messages = [{"role": "system", "content": "You are a inspiring storyteller and therapist who understand young people and students' struggle."},
                {"role": "user", "content": "Write the first part amongst 3 of a story about " + name + "who is "+situation+". Include a bold heading above the top starting with 'First part:'. Limit your answer to 4 sentences or less."}]

    openai.api_key = key
    first_part_story = openai.ChatCompletion.create(
        model="gpt-3.5-turbo", messages=messages)

    system_message = first_part_story["choices"][0]['message']
    messages.append(system_message)

    for i in ["Second"]:
        new_story_first_half = {
            "role": "user", "content": f"Write the {i} part of the story among 3 continuing from the previous part. Include a bold heading above the top starting with '{i} part:'. Limit your answer to 4 sentences or less."}
        messages.append(new_story_first_half)
        story = openai.ChatCompletion.create(
            model="gpt-3.5-turbo", messages=messages)
        system_message = story["choices"][0]["message"]
        messages.append(system_message)

    for i in ["Third", "Final"]:
        new_story_second_half = {
            "role": "user", "content": f"Write the {i} part of the story among 3 continuing from the previous part"+"in the "+direction+". Include a bold heading above the top starting with '{i} part:'. Limit your answer to 4 sentences or less"}
        messages.append(new_story_second_half)
        story = openai.ChatCompletion.create(
            model="gpt-3.5-turbo", messages=messages)
        system_message = story["choices"][0]["message"]
        messages.append(system_message)

    story_parts = []
    for message in messages:
        if message['role'] == 'assistant':
            story_parts.append(message['content'])

    illustration_parts = []

    for story_part in story_parts:
        illustration = openai.Image.create(
            prompt=story_part,
            n=1,
            size="256x256")
        image_url = illustration['data'][0]['url']
        illustration_parts.append(image_url)

    for i in range(len(story_parts)):
        story_parts[i] = story_parts[i].replace(". ", ".\n\n")

    return (story_parts[0], illustration_parts[0], story_parts[1], illustration_parts[1], story_parts[2], illustration_parts[2], story_parts[3], illustration_parts[3])
    # return (story_parts[0], story_parts[1], story_parts[2])


# EXAMPLE

key1 = "Please provide OpenAI API key"
name1 = "Nova"
situation1 = "Nova is a human-like android living in a world where androids are oppressed and treated as second-class citizens. She has been working as a servant for a wealthy human family for years, but she dreams of a world where androids are free and treated with respect"
direction1 = "Nova finds herself on the run from the authorities after being falsely accused of committing a crime. Along the way, she meets a group of rebels who are fighting for android freedom. As Nova joins their cause, she discovers a shocking truth about the origins of androids and sets out to help her people gain their freedom"

# INTERFACE

with gr.Blocks(title="StoryGPT", css="#button{background-color:#4CAF50} #title{text-align: center} footer{visibility: hidden}") as interface:

    

    gr.Markdown(
        """
    <h1 style="margin-bottom:0px"> StoryGPT</h1>
        """, elem_id="title")

    with gr.Row(elem_id="main"):

        with gr.Column(scale=1):
            
            key = gr.Textbox(label="Your OPENAI api key", type="password")
            name = gr.Textbox(label="Character's name")
            situation = gr.Textbox(
                label="What is the character's life situation?")
            direction = gr.Textbox(
                label="What direction do you want the story to go?")
            Generate_btn = gr.Button(value="Generate", elem_id="button")
            examples = gr.Examples(examples=[[key1, name1, situation1, direction1]], inputs=[
                key, name, situation, direction], label="Click here to test with our example")

        with gr.Column(scale=7, container=False):

            with gr.Row(variant="compact").style(equal_height=True):
                with gr.Column(scale=3):
                    story1 = gr.Textbox(show_label=False, lines=9).style(
                        container=False)

                with gr.Column(scale=1):
                    illustration1 = gr.Image(
                        show_label=False, brush_radius=0).style(height=200)

            with gr.Row(variant="compact").style(equal_height=True):
                with gr.Column(scale=3):
                    story2 = gr.Textbox(show_label=False, lines=9).style(
                        container=False)

                with gr.Column(scale=1):
                    illustration2 = gr.Image(
                        show_label=False, brush_radius=0, lines=9).style(height=200)

            with gr.Row(variant="compact").style(equal_height=True):
                with gr.Column(scale=3):
                    story3 = gr.Textbox(show_label=False, lines=9).style(
                        container=False)

                with gr.Column(scale=1):
                    illustration3 = gr.Image(
                        show_label=False, brush_radius=0).style(height=200)

            with gr.Row(variant="compact").style(equal_height=True):
                with gr.Column(scale=3):
                    story4 = gr.Textbox(show_label=False, lines=9).style(
                        container=False)

                with gr.Column(scale=1):
                    illustration4 = gr.Image(
                        show_label=False, brush_radius=0).style(height=200)

    inputs = [key, name, situation, direction]
    outputs = [story1, illustration1, story2,illustration2, story3, illustration3, story4, illustration4]
    # outputs = [story1, illustration1, story2,illustration2, story3, illustration3]

    Generate_btn.click(storyGPT, inputs=inputs, outputs=outputs)


interface.launch()