import os

import carb
import carb.profiler

import omni.ext
import omni.ui

WINDOW_NAME = "Images testing"


class ImagesWindow:
    def _update_byte_image(self):
        import math

        size_increase = max(16.0 * math.sin(self._time), 0.0)

        size_x = (int)(32 + size_increase)
        size_y = (int)(32 + size_increase)

        byte_data = []
        byte_data_float = []
        for x in range(size_x):
            for y in range(size_y):
                color = int(255 * (math.sin(self._time + 0.01 * (x * x + y * y)) * 0.4 + 0.6))
                color_float = math.sin(self._time + 0.01 * (x * x + y * y)) * 0.4 + 0.6
                byte_data.extend([0, color, 0, color])
                byte_data_float.extend([color_float, color_float, 0, color_float])

        self._time += 100.0
        # This will call version that uses byte array, and has a default format omni.ui.TextureFormat.RGBA8_UNORM
        self._byte_provider.set_bytes_data(byte_data, [size_x, size_y])
        # This will call version that uses array of floating point numbers, and has a default format omni.ui.TextureFormat.RGBA32_SFLOAT
        self._byte_provider_float.set_bytes_data(byte_data_float, [size_x, size_y])

    def _update_dynamic_frame(self):
        self._raster_provider_dynamic = omni.ui.RasterImageProvider()
        self._raster_provider_dynamic.source_url = r"${exe-path}/resources/icons/material.png"
        with self._dynamic_frame:
            self._img2_dynamic = omni.ui.ImageWithProvider(self._raster_provider_dynamic)

    def startup(self):
        carb.profiler.begin(1, "Creating explicit image providers")
        self._window = omni.ui.Window(WINDOW_NAME, width=600, height=340)

        self._time = 0.0
        self._byte_provider = omni.ui.ByteImageProvider()
        self._byte_provider_float = omni.ui.ByteImageProvider()
        self._update_byte_image()

        self._raster_provider = omni.ui.RasterImageProvider()
        self._raster_provider.source_url = r"${exe-path}/resources/icons/material.png"

        self._vector_provider = omni.ui.VectorImageProvider()
        self._vector_provider.source_url = r"${exe-path}/resources/glyphs/toolbar_select_prim.svg"

        carb.log_warn("image: %dx%d" % (self._vector_provider.width, self._vector_provider.height))
        carb.profiler.end(1)

        styles = [
            {
                "": {"image_url": "resources/icons/Nav_Walkmode.png"},
                ":hovered": {"image_url": "resources/icons/Nav_Flymode.png"},
            },
            {
                "": {"image_url": "resources/icons/Move_local_64.png"},
                ":hovered": {"image_url": "resources/icons/Move_64.png"},
            },
            {
                "": {"image_url": "resources/icons/Rotate_local_64.png"},
                ":hovered": {"image_url": "resources/icons/Rotate_global.png"},
            },
        ]

        carb.profiler.begin(1, "Building UI")

        def set_image(model, image):
            value = model.get_item_value_model().get_value_as_int()
            image.set_style(styles[value])

        with self._window.frame:
            with omni.ui.VStack():
                with omni.ui.HStack():
                    with omni.ui.HStack():
                        self._img1 = omni.ui.ImageWithProvider(self._byte_provider)
                        self._img1_float = omni.ui.ImageWithProvider(self._byte_provider_float)
                    self._btn1update = omni.ui.Button(text="Update")
                    self._btn1update.set_clicked_fn(self._update_byte_image)
                with omni.ui.HStack():
                    self._img2 = omni.ui.ImageWithProvider(self._raster_provider)
                    self._dynamic_frame = omni.ui.Frame(width=100, height=100)
                    self._btn2refresh = omni.ui.Button(text="Refresh")
                    self._btn2refresh.set_clicked_fn(self._update_dynamic_frame)
                self._img3 = omni.ui.ImageWithProvider(self._vector_provider)
                self._img4 = omni.ui.ImageWithProvider(r"${exe-path}/resources/glyphs/timeline_loop.svg")
                self._img6 = omni.ui.ImageWithProvider(width=64, height=64, style=styles[0])
                with omni.ui.HStack(width=omni.ui.Percent(50)):
                    omni.ui.Label("Select a texture to display", name="text")
                    model = omni.ui.ComboBox(0, "Navigation", "Move", "Rotate").model
                    model.add_item_changed_fn(lambda m, i, im=self._img6: set_image(m, im))
        carb.profiler.end(1)

    def shutdown(self):
        self._img1 = None
        self._img1_float = None
        self._img2 = None
        self._img3 = None
        self._img4 = None
        self._img6 = None

        self._byte_provider = None
        self._raster_provider = None
        self._vector_provider = None

        self._window = None


class Extension(omni.ext.IExt):
    def __init__(self):
        super().__init__()
        pass

    def on_startup(self):
        self._ext_window = ImagesWindow()
        self._ext_window.startup()

    def on_shutdown(self):
        self._ext_window.shutdown()
        self._ext_window = None
