from utils import *


taskType = "3"
temp_examples = get_temps_examples(taskType)
user_examples = get_user_examples(taskType)
showcase_examples = get_showcase_examples(taskType)
user_recorder = UserRecorder()

# 当试用次数耗尽，显示的信息
no_more_attempts = "Sorry, You've used all trial attempts. Please enter your token to continue..."

# Description
title = r"""
<h1 align="center">Selfit AI App: Product poster background generation</h1>
"""

description = r"""
<a href='https://heybeauty.ai/extension' target='_blank'><b> Go to HeyBeauty for Faster and Free AI App! 🤗 </b></a>.<br>

<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/selfitcamera/Outfit-Anyone-in-the-Wild' target='_blank'><b>Product poster background generation</b></a>.<br>
1. Please wait for a while after submitting the task, you will get the result within 30 seconds.
2. Share app with your friends and enjoy! 😊.<br>
"""

css = """
.gradio-container {width: 85% !important}
"""


def onClick(temp_image, user_image, caption_text, request: gr.Request):
    print("======> temp_image ", type(temp_image), temp_image)
    print("======> user_image ", type(user_image))
    print("======> caption_text ", type(caption_text))
    if temp_image is None: 
        yield None, "please choose a template background!!!"
        return None, "please choose a template background!!!"

    try:
        client_ip = request.client.host
        x_forwarded_for = dict(request.headers).get('x-forwarded-for')
        if x_forwarded_for: client_ip = x_forwarded_for
        if not check_region_warp(client_ip):
            return None, "Failed !!! Our server is under maintenance, please try again later"
            
        _, total_n, _ = user_recorder.get_record(client_ip)
        if total_n>=LimitTask:
            print(no_more_attempts)
            yield None, no_more_attempts
            return None, "no_more_attempts"

        upload_url = upload_user_img(client_ip, user_image)
        if len(upload_url)==0:
            yield None, "fail to upload"
            return None, "fail to upload"

        taskId = publicSelfitTask(upload_url, temp_image, caption_text)
        if not taskId:
            yield None, "fail to public task..."
            return None, "fail to public task..."

        # taskId = "95a19e33-1ff5-446e-ac37-d3a8dd59eba8-e1"
        max_try = 30
        wait_s = 3
        time.sleep(5)
        for i in range(max_try):
            time.sleep(wait_s)
            taskStatus = getTaskRes(taskId)
            if taskStatus is None: continue
            user_recorder.save_record(taskStatus, ip=client_ip)

            status = taskStatus['status']
            if status in ['FAILED', 'CANCELLED', 'TIMED_OUT', ]:
                yield None, f"task failed, query {i}, status {status}"
                return None, f"task failed, query {i}, status {status}"
            elif status in ['IN_QUEUE', 'IN_PROGRESS', 'IN_QUEUE', ]:
                yield None, f"task is on processing, query {i},  status {status}"
            elif status=='COMPLETED':
                out = taskStatus['output']['job_results']['output1']
                yield out, f"task is COMPLETED"
                return out, f"{i} task COMPLETED"
        yield None, "fail to query task.."
        return None, "fail to query task.."
    except Exception as e:
        print(e)
        # raise e
        yield None, "fail to create task"
        return None, "fail to create task"

def onLoad(request: gr.Request):
    client_ip = request.client.host
    x_forwarded_for = dict(request.headers).get('x-forwarded-for')
    if x_forwarded_for:
        client_ip = x_forwarded_for
    his_datas, total_n, msg = user_recorder.get_record(client_ip)
    left_n = max(0, LimitTask-total_n)
    his_datas.append(msg)
    his_datas.append(f"Submit ({left_n} attempts left)")
    return his_datas

with gr.Blocks(css=css) as demo:
    gr.Markdown(title)
    gr.Markdown(description)

    with gr.Row():
        with gr.Column():
            with gr.Column():
                temp_image = gr.Image(sources='clipboard', type="filepath", label="Choose a template", 
                    value=temp_examples[0][0])
                temp_example = gr.Examples(inputs=[temp_image],
                                      examples_per_page=9, examples=temp_examples)
        with gr.Column():
            with gr.Column():
                user_image = gr.Image(value=None, type="numpy", label="product photo")
                caption_text = gr.Textbox(value="", interactive=True, 
                    label='English caption text(optional)') 
        with gr.Column():
            with gr.Column():
                res_image = gr.Image(label="generate image", value=None, type="filepath")
                info_text = gr.Textbox(value="", interactive=False, 
                    label='runtime information')
                run_button = gr.Button(value="Submit")
                MK01 = gr.Markdown()
            
    with gr.Column():
        show_case = gr.Examples(
            examples=showcase_examples,
            inputs=[temp_image, user_image, res_image, ],
            label=None
        )

    with gr.Tab('history'):

        with gr.Row():
            MK02 = gr.Markdown()

        with gr.Row():
            his_input1 = gr.HTML()
            his_output1 = gr.HTML()

        with gr.Row():
            his_input2 = gr.HTML()
            his_output2 = gr.HTML()

        with gr.Row():
            his_input3 = gr.HTML()
            his_output3 = gr.HTML()            

    run_button.click(fn=onClick, inputs=[temp_image, user_image, caption_text], 
                     outputs=[res_image, info_text], concurrency_limit=50)
    demo.load(onLoad, inputs=[], outputs=[his_input1, his_output1,
        his_input2, his_output2, his_input3, his_output3,
        MK02, run_button])

if __name__ == "__main__":

    demo.queue(max_size=50)
    # demo.queue(concurrency_count=60)
    # demo.launch(server_name='0.0.0.0', server_port=225)
    demo.launch(server_name='0.0.0.0')