# === Included SDK === #
import gradio as gr

# === Installed from requirements.txt === #
from PIL import Image
import numpy as np
import torch

# === Python built-in === #
import sys
import os

# ============================================================ #
# Q: Why use gr.Blocks instead of gr.Interface?                #
# A: Only Blocks supports Layouts (and I'm more used to this)  #
# ============================================================ #
with gr.Blocks() as demo:

    # ======================================== #
    # Markdown is pointless in API usage... XD #
    # ======================================== #
    gr.Markdown("""<h1 align="center">Hello World</h1>""")
    gr.Markdown("""<p align="center">For <b>learning</b> purposes only...</p>""")

    # ====================================== #
    # function name is overrides by api_name #
    # ====================================== #
    def log() -> str:
        info: list[str] = []
        info.append(f"OS: {os.name}")
        info.append(f"Python: {sys.version.split(' ', 1)[0]}")
        info.append(f"CUDA: {torch.cuda.is_available()}")

        return "\n".join(info)

    def proc_image(
        img: np.ndarray, target_w: int, target_h: int
    ) -> list[str, Image.Image]:

        input_h, input_w, channels = img.shape

        log = f"Resize a {input_w}x{input_h} image to {target_w}x{target_h}..."

        image = Image.fromarray(img)
        image = image.resize((target_w, target_h))

        return [log, image]

    with gr.Row():

        with gr.Column():

            # ============================= #
            # Images take base64 str in API #
            # ============================= #

            input_image = gr.Image(
                value=None,
                image_mode="RGB",
                sources="upload",
                type="numpy",
                label="Input",
                show_label=True,
                show_download_button=False,
                visible=True,
                interactive=True,
                show_share_button=False,
                show_fullscreen_button=False,
            )

            with gr.Row():
                target_w = gr.Slider(
                    minimum=1,
                    maximum=1024,
                    step=1,
                    value=512,
                    label="width",
                    show_label=True,
                    interactive=True,
                )

                target_h = gr.Slider(
                    minimum=1,
                    maximum=1024,
                    step=1,
                    value=512,
                    label="width",
                    show_label=True,
                    interactive=True,
                )

            output_image = gr.Image(
                value=None,
                image_mode="RGB",
                sources="upload",
                type="pil",
                label="Output",
                show_label=True,
                show_download_button=False,
                visible=True,
                interactive=False,
                show_share_button=False,
                show_fullscreen_button=False,
            )

            btn1 = gr.Button("Resize")

        with gr.Column():

            # ============================= #
            # Buttons do not show up in API #
            # ============================= #

            btn2 = gr.Button("Print Info")

    # ==================================== #
    # label is shown in the returns in API #
    # ==================================== #

    console_logs = gr.Text(
        value="...",
        label="Logs",
        show_label=True,
        lines=1,
        max_lines=8,
        interactive=False,
    )

    # =================================== #
    # api_name is used as endpoint in API #
    # =================================== #

    btn1.click(
        fn=proc_image,
        inputs=[input_image, target_w, target_h],
        outputs=[console_logs, output_image],
        show_progress="hidden",
        show_api=True,
        api_name="resizeImage",
    )

    btn2.click(
        fn=log,
        inputs=None,
        outputs=console_logs,
        show_progress="hidden",
        show_api=True,
        api_name="systemInfo",
    )


demo.launch()