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

HF_TOKEN = os.environ.get('HF_KEY')

client = Client("leaf1170124460/mattingPrivate",
                max_workers=3,
                hf_token=HF_TOKEN)

def process_image_check(path_input):
    if path_input is None:
        raise gr.Error(
            "Missing image in the left pane: please upload an image first."
        )

def process_pipe_matting(matting_image_input):
    return client.predict(
      path_input=handle_file(matting_image_input),
      api_name="/gradio_handler"
    )

def run_demo_server():
    gradio_theme = gr.themes.Default()
    with gr.Blocks(
        theme=gradio_theme,
        title="Matting",
    ) as demo:
        with gr.Row():
            gr.Markdown("# Matting Demo")
        with gr.Row():
            gr.Markdown("### Due to the GPU quota limit, if an error occurs, please wait for 5 minutes before retrying.")
        with gr.Row():
            with gr.Column():
                matting_image_input = gr.Image(
                    label="Input Image",
                    type="filepath",
                )
                with gr.Row():
                    matting_image_submit_btn = gr.Button(
                        value="Estimate Matting", variant="primary"
                    )
                    matting_image_reset_btn = gr.Button(value="Reset")
            with gr.Column():
                matting_image_output = gr.Image(label='Output')


        matting_image_submit_btn.click(
            fn=process_image_check,
            inputs=matting_image_input,
            outputs=None,
            preprocess=False,
            queue=False,
        ).success(
            fn=process_pipe_matting,
            inputs=[
                matting_image_input,
            ],
            outputs=[matting_image_output],
            concurrency_limit=1,
        )

        matting_image_reset_btn.click(
            fn=lambda: (
                None,
                None,
            ),
            inputs=[],
            outputs=[
                matting_image_input,
                matting_image_output,
            ],
            queue=False,
        )

    demo.queue(
        api_open=False,
    ).launch()


if __name__ == '__main__':
    run_demo_server()