# === Chosen SDK === #
import gradio as gr

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

# === Python built-in === #
from io import BytesIO
import base64
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 overridden by api_name #
    # ======================================= #
    def whatever() -> 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()}")
        info.append(f"Open-CV: {cv2.__version__}")
        info.append(f"torch: {getattr(torch, 'float16', False)}")
        info.append(f"np: {getattr(np, 'float16', False)}")

        return "\n".join(info)

    # ===================================================== #
    # parameter name is also shown as parameter name in API #
    # ===================================================== #
    def proc_image(image: str, width: int, height: int) -> list[str, Image.Image]:
        assert isinstance(image, str)

        # ==================================== #
        # This works with base64 image via API #
        # ==================================== #
        image = Image.open(BytesIO(base64.b64decode(image)))

        if isinstance(width, float):
            width = int(width)
        if isinstance(height, float):
            height = int(height)

        input_w, input_h = image.size

        log = f"Resize a {input_w}x{input_h} image to {width}x{height}..."

        image = image.convert("RGB").resize((width, height))

        with BytesIO() as buffer:
            image.save(buffer, format="PNG")
            b64_image = base64.b64encode(buffer.getvalue()).decode()

        return [log, b64_image]

    with gr.Row():

        with gr.Column():

            # ==================================== #
            # variable name does not matter in API #
            # ==================================== #
            foo1 = gr.Text(
                value="",
                label="Input Image",
                show_label=True,
                info="Image in base64",
                visible=True,
                interactive=True,
                lines=1,
                max_lines=1,
            )

            # =============================================== #
            # Only value matters (shown as the Default value) #
            # =============================================== #
            with gr.Row():
                foo2 = gr.Slider(
                    minimum=1,
                    maximum=1024,
                    step=1,
                    value=512,
                    label="width",
                    show_label=True,
                    interactive=True,
                )

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

            bar = gr.Text(
                value="",
                label="Output Image",
                show_label=True,
                info="Image in base64",
                visible=True,
                interactive=False,
                lines=1,
                max_lines=1,
            )

            yeet1 = gr.Button("Resize")

        with gr.Column():

            # ============================== #
            # Button does not show up in API #
            # ============================== #
            yeet2 = 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 #
    # =================================== #
    yeet2.click(
        fn=whatever,
        inputs=None,
        outputs=console_logs,
        show_progress="hidden",
        show_api=True,
        api_name="systemInfo",
    )

    # =============================================== #
    # API order is based on event declaration in code #
    # =============================================== #
    yeet1.click(
        fn=proc_image,
        inputs=[foo1, foo2, foo3],
        outputs=[console_logs, bar],
        show_progress="hidden",
        show_api=True,
        api_name="resizeImage",
    )

demo.launch(show_error=True)