import gradio as gr

from fawkes.protection import Fawkes


def run_protection_interface(uploaded_image, mode='low', sd=1e6, format='png', separate_target=False, no_align=False,
                             debug=False):
    """
    Gradio compatible function for running protection.
    """
    if uploaded_image is None:
        return None, "No image uploaded."

    # Run the protection process
    protector = Fawkes(gpu="0", batch_size=1, mode=mode)
    processed_image = protector.run_protection(uploaded_image, sd=sd, batch_size=1, format=format,
                                               separate_target=separate_target, debug=debug, no_align=no_align)

    if processed_image is not None:
        return processed_image, "Protection process completed."
    else:
        return None, "Protection process failed or no cloaked image generated."


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Upload Image")
            uploaded_image = gr.Image(type="pil", label="Upload Image")

        with gr.Column():
            gr.Markdown("### Configuration Options")
            mode = gr.Radio(label="Mode", choices=['low', 'mid', 'high'], value='low')
            format = gr.Radio(label="Output Format", choices=['png', 'jpg', 'jpeg'], value='png')
            separate_target = gr.Checkbox(label="Separate Target")
            no_align = gr.Checkbox(label="No Align")
            with gr.Accordion(label='Advanced Config', open=False):
                sd = gr.Slider(label="Penalty Number (SD)", minimum=1e5, maximum=1e7, value=1e6)

    run_button = gr.Button("Run Protection")
    output_image = gr.Image(label="Processed Image")
    output_text = gr.Textbox(label="Output Message")

    run_button.click(
            fn=run_protection_interface,
            inputs=[uploaded_image, mode, sd, format, separate_target, no_align],
            outputs=[output_image, output_text]
            )

if __name__ == "__main__":
    demo.launch()