import omni.ext
import omni.kit.commands
import omni.kit.ui
import omni.kit.debug.python



EXTENSION_NAME = "VS Code Link"



class DebugBreak(omni.kit.commands.Command):
    def do(self):
        omni.kit.debug.python.breakpoint()


def toColor01(color):
    return tuple(c / 255.0 for c in color) + (1,)


class Extension(omni.ext.IExt):
    """
    Extension to enable connection of VS Code python debugger. It enables it using python ptvsd module and shows status.
    """

    def on_startup(self):
        self._window = omni.kit.ui.Window(EXTENSION_NAME, 300, 150)
        layout = self._window.layout
        row_layout = layout.add_child(omni.kit.ui.RowLayout())

        # Status label
        row_layout.add_child(omni.kit.ui.Label("Status:"))
        self._status_label = row_layout.add_child(omni.kit.ui.Label())

        wait_info = omni.kit.debug.python.get_listen_address()
        info_label = layout.add_child(omni.kit.ui.Label())
        info_label.text = f"Address: {wait_info}" if wait_info else "Error"

        # Break button
        self._break_button = layout.add_child(omni.kit.ui.Button("Break"))

        def on_click(*x):
            omni.kit.commands.execute("DebugBreak")

        self._break_button.set_clicked_fn(on_click)

        # Monitor for attachment
        self._attached = None
        self._set_attached(False)

        def on_update(dt):
            self._set_attached(omni.kit.debug.python.is_attached())

        self._window.set_update_fn(on_update)

    def _set_attached(self, attached):
        if self._attached != attached:
            ATTACH_DISPLAY = {
                True: ("VS Code Debugger Attached", (0, 255, 255)),
                False: ("VS Code Debugger Unattached", (255, 0, 0)),
            }
            self._attached = attached
            self._break_button.enabled = attached
            self._status_label.text = ATTACH_DISPLAY[self._attached][0]
            self._status_label.set_text_color(toColor01(ATTACH_DISPLAY[self._attached][1]))
