import gradio as gr
import random
import time
import requests
import io
from PIL import Image
import traceback

from base64 import b64decode,b64encode
from io import BytesIO
from better_profanity import profanity


    
with gr.Blocks(theme="darkdefault") as demo:

    def welcome(name):
        return f"Welcome to AIXRPL.com Minter, {name}!"




    def profanityCheck(prompt):
            prompt = prompt.replace('+',' ').replace('|',' ')
            if profanity.contains_profanity(prompt):
                return True
            else:
                return False


    def inference(_prompt,_token):
        try:
            from PIL import Image
            import uuid
            import os
            print(_prompt,_token)
            
            if profanityCheck(_prompt):
                    img = Image.open('unsafe.png') 
                    return img,'unsafe','','',''        
            
            r = requests.post(url='https://aixrplart-5czkww5hsa-uc.a.run.app/create',data={"prompt":_prompt,"token":_token})
            all_data = r.json()
            print(all_data.keys())
            
            import base64
            from io import BytesIO
            from PIL import Image
            
            im_bytes = base64.b64decode(all_data['img_data'])   # im_bytes is a binary image
            im_file = BytesIO(im_bytes)  # convert image to file-like object
            img = Image.open(im_file)   # img is now PIL Image object
            
            
            return(img,all_data['description'],all_data['image_url'],all_data['keywords'],all_data['keywords_string'])
        except Exception as e:
                    print('exception:',e)
                    traceback.print_exc()
                    return '','','','',''
        # img.save('/tmp/data.png')
        #return '/tmp/data.png'


    with gr.Group():                
        generate_progress = gr.StatusTracker(cover_container=True)

        with gr.Row():
            with gr.Column():
              with gr.Tab("Create"):
                  
                gr.Markdown(
                    """
                    Create AI generated artworks by using prompt engineering.
                    """
                        )

                text = gr.Textbox(
                    label="Enter Prompt", show_label=True, max_lines=5
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=True,
                )

                btn = gr.Button("Create").style(
                    margin=True,
                    rounded=(False, True, True, False),
                )

                gr.Markdown(
                    """
                    AI generated metadata.
                    """
                        )

                description = gr.Textbox(
                    label="AI Generated Description", interactive=True, show_label=True, max_lines=1, elem_id="descData"
                ).style(
                    border=(True, False, True, True),
                    rounded=(True, False, False, True),
                    container=True,
                )
                
                traits = gr.HighlightedText(label="Auto Traits",interactive=True, show_label=True)
                # build_result = gr.Gallery()#gr.Image(interactive=False, shape=(320,320))



            with gr.Column():
              with gr.Tab("Artwork"):

                        build_result = gr.Image(type="pil", shape=(512,None),show_label=True,label="Artwork Preview",interactive=False,)

                        walletToken = gr.Textbox(
                            visible=False, interactive=True, elem_id="walletToken", max_lines=1
                        )  

                        imageData = gr.Textbox(
                            visible=False, interactive=False, elem_id="imageData", max_lines=1
                        )  

                        attribData = gr.Textbox(
                            visible=False, interactive=False, elem_id="attribData", max_lines=1
                        )  


            btn.click(
                      inference,
                      inputs=[text,walletToken],
                      outputs=[build_result,description,imageData, traits, attribData],
                      status_tracker=generate_progress,
                      api_name="generate"
                  )


if __name__ == "__main__":
    demo.launch(show_api=False, debug=True, enable_queue=True)