import glob
import os.path
import tempfile

import gradio as gr
from PIL import Image

from attack import Attacker, make_args


def attack_given_image(image: Image.Image, target: str, steps: int, eps: float, step_size: float,
                       progress=gr.Progress()):
    if image.mode != 'RGB':
        image = image.convert('RGB')

    with tempfile.TemporaryDirectory() as td_input, tempfile.TemporaryDirectory() as td_output:
        image_filename = os.path.join(td_input, 'image.png')
        image.save(image_filename)

        def _step_func(current_step: int):
            progress(current_step / steps)

        args = make_args([
            image_filename,
            '--out_dir', str(td_output),
            '--target', target,
            '--eps', str(eps),
            '--step_size', str(step_size),
            '--steps', str(steps),
        ])
        attacker = Attacker(args)
        before_prediction = attacker.image_predict(image)
        attacker.attack(args.inputs, _step_func)

        output_filename, *_ = glob.glob(os.path.join(td_output, '*.png'))
        output_image = Image.open(output_filename)
        after_prediction = attacker.image_predict(output_image)

        return before_prediction, after_prediction, output_image


if __name__ == '__main__':
    with gr.Blocks() as demo:
        with gr.Row():
            gr.HTML('<p style=\'font-size: 18px; text-align: center;\'>'
                    '   Too slow? Try this locally: '
                    '   <a href="https://github.com/7eu7d7/anime-ai-detect-fucker" '
                    '      style=\'text-decoration: underline;\'>'
                    '       Github - 7eu7d7/anime-ai-detect-fucker'
                    '   </a>.'
                    '   or'
                    '   <a href="https://colab.research.google.com/drive/1bzoZtE28Y8vfmjuaxXKdh-fr4K0WWSmf?usp=sharing" '
                    '      target="_parent">'
                    '       <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>'
                    '   </a>'
                    '</p>')
        with gr.Row():
            with gr.Column():
                gr_input_image = gr.Image(type='pil', label='Original Image')
                with gr.Row():
                    gr_attack_target = gr.Radio(['auto', 'ai', 'human', 'same'], value='auto', label='Attack Target')
                    gr_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label='Steps')
                with gr.Row():
                    gr_eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
                    gr_noise_step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001,
                                                   value=0.136)
                gr_btn_submit = gr.Button(value='Attack This Image', variant='primary')

            with gr.Column():
                with gr.Row():
                    gr_before_prediction = gr.Label(label='Before Prediction')
                    gr_after_prediction = gr.Label(label='After Prediction')
                gr_output_image = gr.Image(type='pil', label='Attacked Image')

        gr_btn_submit.click(
            attack_given_image,
            inputs=[gr_input_image, gr_attack_target, gr_steps, gr_eps, gr_noise_step_size],
            outputs=[gr_before_prediction, gr_after_prediction, gr_output_image],
        )

    demo.queue(os.cpu_count()).launch()