import omni.ext
from omni.rtx.window.settings import RendererSettingsFactory

from .widgets.common_widgets import CommonSettingStack
from .widgets.post_widgets import PostSettingStack
from .widgets.pt_widgets import PTSettingStack
from .widgets.rt_widgets import RTSettingStack


class RTXSettingsExtension(omni.ext.IExt):

    MENU_PATH = "Rendering"

    rendererNames = ["Real-Time", "Interactive (Path Tracing)"]
    stackNames = ["Common", "Ray Tracing", "Path Tracing", "Post Processing"]

    def on_startup(self):
        RendererSettingsFactory.register_stack(self.stackNames[0], CommonSettingStack)
        RendererSettingsFactory.register_stack(self.stackNames[1], RTSettingStack)
        RendererSettingsFactory.register_stack(self.stackNames[2], PTSettingStack)
        RendererSettingsFactory.register_stack(self.stackNames[3], PostSettingStack)
        RendererSettingsFactory.register_renderer(
            self.rendererNames[0], [self.stackNames[0], self.stackNames[1], self.stackNames[3]]
        )
        RendererSettingsFactory.register_renderer(
            self.rendererNames[1], [self.stackNames[0], self.stackNames[2], self.stackNames[3]]
        )
        RendererSettingsFactory.build_ui()

    def on_shutdown(self):
        for renderer in self.rendererNames:
            RendererSettingsFactory.unregister_renderer(renderer)
        for stack in self.stackNames:
            RendererSettingsFactory.unregister_stack(stack)
        RendererSettingsFactory.build_ui()
