import os
import gradio as gr
from gradio_client import Client, handle_file
from PIL import Image

def predict(imgs, garm_img):
    print(imgs, garm_img)
    client = Client("life4cut/ff-v1", hf_token=os.environ.get('HF_TOKEN_FF'))
    result = client.predict(
        #dict={"background":handle_file('https://s3.ap-northeast-2.amazonaws.com/life4cut.app/images/temp/111.png'),"layers":[],"composite":None},
        #garm_img=handle_file('https://s3.ap-northeast-2.amazonaws.com/life4cut.app/images/temp/00_style5.png'),
        dict={"background":handle_file(imgs),"layers":[],"composite":None},
        garm_img=handle_file(garm_img),
        garment_des="Hello!!",
        is_checked=True,
        is_checked_crop=False,
        denoise_steps=30,
        seed=42,
        api_name="/tryon"
    )
    #print(result)
    return result[0], result[1]
#print(result[1])

# View the image
#Image.open(result[0])
#Image.open(result)
example_path = os.path.join(os.path.dirname(__file__), 'example')
garm_list = os.listdir(os.path.join(example_path,"cloth"))
garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]
human_list = os.listdir(os.path.join(example_path,"human"))
human_list_path = [os.path.join(example_path,"human",human) for human in human_list]

image_blocks = gr.Blocks().queue()
with image_blocks as demo:
    gr.Markdown("## fashion filter")
    with gr.Row():
        with gr.Column():
            #imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
            imgs = gr.Image(sources='upload', type="filepath", label='Human. Mask with pen or use auto-masking')
            with gr.Row():
              is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
            with gr.Row():
              is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
            example = gr.Examples(
                inputs=imgs,
                examples_per_page=10,
                examples=human_list_path
            )
        with gr.Column():
            garm_img = gr.Image(label="Garment", sources='upload', type="filepath")
            example = gr.Examples(
                inputs=garm_img,
                examples_per_page=10,
                examples=garm_list_path)          
        with gr.Column():
            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
            masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
        with gr.Column():
            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
            image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False) 
    with gr.Column():
        try_button = gr.Button(value="predict")
    try_button.click(fn=predict, inputs=[imgs, garm_img], outputs=[image_out,masked_img])        
image_blocks.launch()