import gradio as gr
from utils import *

with gr.Blocks() as demo:
    gr.Markdown("# Encrypt the secret message into image.")
    with gr.Tab("Encrypt"):
        with gr.Row():
            with gr.Column():
                encrypt_msg = gr.Textbox(lines=1, label="Encrypt Message")
                encrypt_key = gr.Textbox(lines=1, label="Encrypt Key")
                encrypt_image = gr.Image()
            encrypt_output = gr.Image()
        encrypt_button = gr.Button("Encrypt")
    with gr.Tab("Decrypt"):
        with gr.Row():
            with gr.Column():
                decrypt_key = gr.Textbox(lines=1, label="Decrypt Key")
                decrypt_image = gr.Image()
            decrypt_output = gr.Textbox(lines=1, label="Decrypt Message")
        decrypt_button = gr.Button("Decrypt")

    encrypt_button.click(encrypt, inputs=[encrypt_msg, encrypt_key, encrypt_image], outputs=[encrypt_output])
    decrypt_button.click(decrypt, inputs=[decrypt_key, decrypt_image], outputs=[decrypt_output])

demo.launch();