import gradio as gr

def dmini(prompt):
    from craiyon import Craiyon
    import time
    from datetime import datetime
    import base64
    from io import BytesIO
    from PIL import Image
    import re

    # if prompt is empty, log to console an error
    if prompt == "":
        return ["error.png"] * 9
    # get time now gmt
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    # format it to a string
    timestamp = str(timestamp)
    # convert string prompt to lowercase
    prompt = prompt.lower()
    # create list bad_words_list from bad_words.txt and strip whitespace
    with open('bad_words.txt') as f:
        bad_words_list = f.readlines()
        bad_words_list = [x.strip() for x in bad_words_list]
        # check if bad_words_list is in prompt
        if any(re.findall(r'\b{}\b'.format(bad_word), prompt) for bad_word in bad_words_list):
            # check if any bad words are in prompt
            blocked_words = []
            for word in bad_words_list:
                if word in prompt:
                    blocked_words.append(word)   
            blocked_words = ', '.join(blocked_words)
            print(timestamp + ": BLOCKED PROMPT \"" + prompt + "\" BANNED WORDS: " + blocked_words)  
            time.sleep(8)
            print("User has had their time wasted.")
            return ["unsafe.png"] * 9
        else:
            generator = Craiyon() # Instantiates the api wrapper
            result = generator.generate(prompt) # Generates 9 images by default and you cannot change that
            images = result.images # A list containing image data as base64 encoded strings
            imagelist = []
            for i in images:
                image = Image.open(BytesIO(base64.decodebytes(i.encode("utf-8"))))
                imagelist.append(image)
            print(timestamp + ": PROMPT \"" + prompt + "\" WAS GENERATED SUCCESSFULLY\n")
            return imagelist

examples = ["Futuristic utopian city in Coruscant. god beams. white buildings. street view. golden sunset. colorful, cityscape, green trees, future, urban, megapolis, breathtaking, Imaginative, utopia, Wallpaper, beautiful design, scifi, high detail, global illumination, ArtStation, art by Sebastian Wagner, Stephan Martiniere, Leon Tukker, Vitaly-Sokol",
"A dream of a distant galaxy, by Caspar David Friedrich, matte painting trending on artstation HQ",
"Cute close portrait dainty beautiful futurstic white robotic girl, portraits, white cyberpunk inflatable formfitting tops, highly detailed sharp big eyes, white fringeshort hair, transparent intricate details, professional 3d visualisation in pastel colours, by wlop, intricate linework, trending on artstation, unreal engine 5 highly rendered, epic composition, by guweiz, shal. e, laica chrose, final fantasy",
"1990s movie, orbit space new york city street with many pedestrians bejng attacked by ufos as a loading screen, cinestill 800t 18mm, heavy grainy picture, very detailed, high quality, 4k panoramic, dramatic lightning, streetlight at night, rain, mud, foggy, anchray flags",
"Portrait of a woman by Greg Rutkowski, she is about 30 years old, pretty, blond hair with two strans around her face, slavic features, melancholic gaze, pretty aquiline nose, she is wearing a blue utilitarian jumpsuit, highly detailed portrait, digital painting, artstation, concept art, smooth, sharp foccus ilustration, Artstation HQ.",
"A small cabin on top of a snowy mountain in the style of Disney, artstation"]


with gr.Blocks() as demo:
    output = gr.Gallery(label="Image Generation", elem_id="output").style(grid=3)
    name = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate. Longer and more detailed prompts work better.", elem_id="name") 
    greet_btn = gr.Button("Generate Image", elem_id="greet_btn")
    ex = gr.Examples(examples=examples, fn=dmini, inputs=name, outputs=output, cache_examples=True)
    ex.dataset.headers = [""]
    greet_btn.click(fn=dmini, inputs=name, outputs=output)

demo.queue(concurrency_count=20, max_size=20).launch(max_threads=150)