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, 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', '0.135914',
            '--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__':
    interface = gr.Interface(
        attack_given_image,
        inputs=[
            gr.Image(type='pil', label='Original Image'),
            gr.Radio(['auto', 'ai', 'human'], value='auto', label='Attack Target'),
            gr.Slider(minimum=1, maximum=50, value=20, step=1, label='Steps'),
            gr.Slider(minimum=1.0, maximum=16.0, value=1.0, step=1 / 8, label='Eps'),
        ],
        outputs=[
            gr.Label(label='Before Prediction'),
            gr.Label(label='After Prediction'),
            gr.Image(type='pil', label='Attacked Image'),
        ],
        interpretation="default"
    )
    interface.queue(os.cpu_count()).launch()