import gradio as gr 
from steganography import Steganography
from util_functions.utils import draw_multiple_line_text, generate_qr_code
from SSL_watermark import encode, decode


TITLE = """<h2 align="center"> ✍️ Invisible Watermark </h2>"""


def apply_watermark(radio_button, input_image, watermark_image, watermark_text, watermark_url):
    input_image = input_image.convert('RGB')
    print(f'radio_button: {radio_button}')
    if radio_button == "Image":
        watermark_image = watermark_image.resize((input_image.width, input_image.height)).convert('L').convert('RGB')
        return Steganography().merge(input_image, watermark_image, digit=7)
    elif radio_button == "Text":
        watermark_image = draw_multiple_line_text(input_image.size, watermark_text)
        return Steganography().merge(input_image, watermark_image, digit=7)
    elif radio_button == "QRCode":
        size = min(input_image.width, input_image.height)
        watermark_image = generate_qr_code(watermark_url).resize((size, size)).convert('RGB')
        return Steganography().merge(input_image, watermark_image, digit=7)
    else:
        print('start encoding ssl watermark...')
        return encode(input_image, epochs=5)

def extract_watermark(extract_radio_button, input_image_to_extract):
    if extract_radio_button == 'Steganography':
        return Steganography().unmerge(input_image_to_extract.convert('RGB'), digit=7).convert('RGBA')
    else:
        decoded_info = decode(image=input_image_to_extract)
        return draw_multiple_line_text(input_image_size=input_image_to_extract.size, text=decoded_info)


with gr.Blocks() as demo:
    gr.HTML(TITLE)
    with gr.Tab("Add watermark"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Image to apply watermark")
                input_image = gr.Image(type='pil')
                with gr.Blocks():
                    gr.Markdown("### Which type of watermark you want to apply?")
                    radio_button = gr.Radio(
                        choices=["QRCode", "Text", "Image", "SSL Watermark"], 
                        label="Watermark type", 
                        value="QRCode",
                        # info="Which type of watermark you want to apply?"
                    )
                    watermark_url = gr.Textbox(
                        placeholder="URL to generate QR code", 
                        visible=True
                    )
                    watermark_text = gr.Textbox(
                        placeholder="What text you want to use as watermark?",
                        visible=False
                    )
                    watermark_image = gr.Image(
                        type='pil', 
                        visible=False
                    )
                    
                    def update_visibility(radio_value):
                        return { 
                            watermark_image: 
                                { 
                                    "visible":radio_value == "Image", 
                                    "__type__": "update"
                                }, 
                            watermark_text: 
                                {
                                    "visible":radio_value == "Text", 
                                    "__type__": "update"
                                },
                            watermark_url: 
                                {
                                    "visible":radio_value == "QRCode", 
                                    "__type__": "update"
                                }
                            }

            with gr.Column():
                gr.Markdown("### Applied watermark image")
                output_image = gr.Image(show_label=False).style(height=500)
        with gr.Row():
            apply_button =gr.Button("Apply")
        
    with gr.Tab("Extract watermark"):
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Image to extract watermark")
                input_image_to_extract = gr.Image(type='pil')
                extract_radio_button = gr.Radio(
                    choices=["Steganography", "SSL Watermark"], 
                    label="Extract methods", 
                    value="Steganography"
                )
            with gr.Column():
                gr.Markdown("### Extracted watermark")
                extracted_watermark = gr.Image(type='pil').style(height=500)
        extract_button = gr.Button("Extract")
    
    radio_button.change(
        fn=update_visibility, 
        inputs=radio_button,
        outputs=[watermark_image, watermark_text, watermark_url]
    )
    apply_button.click(
        fn=apply_watermark, 
        inputs=[radio_button, input_image, watermark_image, watermark_text, watermark_url], 
        outputs=[output_image]
    )
    extract_button.click(
        fn=extract_watermark, 
        inputs=[extract_radio_button, input_image_to_extract], 
        outputs=[extracted_watermark]
    )

demo.launch()