import asyncio

import omni.ui
import omni.kit.app
import omni.kit.commands
import carb.settings


### Copied from omni.phyx; Kit team will coalesce into omni.ui.settings module
class Model(omni.ui.AbstractValueModel):
    def __init__(self, path):
        super().__init__()
        self._path = path
        val = carb.settings.get_settings().get(path)
        if val is None:
            val = 0
        self._value = val

        def on_change(item, event_type, path=path):
            self._value = carb.settings.get_settings().get(path)
            self._value_changed()

        self._sub = omni.kit.app.SettingChangeSubscription(path, on_change)

    def get_value_as_bool(self):
        return self._value

    def get_value_as_int(self):
        return self._value

    def get_value_as_float(self):
        return self._value

    def get_value_as_string(self):
        return self._value

    def set_value(self, value):
        omni.kit.commands.execute("ChangeSetting", path=self._path, value=value)
        self._value = value
        self._value_changed()

def create_setting_widget(setting_path, cls, **kwargs):
    model = Model(setting_path)
    widget = cls(model, **kwargs)
    return (model, widget)

### end omni.physx boilerplate

class OmniHydraSettings(omni.ui.Window):
    def __init__(self):
        self._title = "OmniHydra Developer Settings"
        super().__init__(self._title, omni.ui.DockPreference.LEFT_BOTTOM, width=800, height=600)
        self._widgets = []

        with self.frame:
            self._build_ui()

            asyncio.ensure_future(self._dock_window(self._title, omni.ui.DockPosition.SAME))

    async def _dock_window(self, window_title: str, position: omni.ui.DockPosition, ratio: float = 1.0):
        frames = 3
        while frames > 0:
            if omni.ui.Workspace.get_window(window_title):
                break
            frames = frames - 1
            await omni.kit.app.get_app().next_update_async()

        window = omni.ui.Workspace.get_window(window_title)
        dockspace = omni.ui.Workspace.get_window("Property")

        if window and dockspace:
            window.dock_in(dockspace, position, ratio=ratio)
            window.dock_tab_bar_visible = False

    def on_shutdown(self):
        self._widgets = []

    def _add_setting_widget(self, name, path, cls, **kwargs):
        with omni.ui.VStack():
            omni.ui.Spacer(height=2)
            with omni.ui.HStack(height=20):
                self._widgets.append(create_setting_widget(path, cls, **kwargs))
                omni.ui.Label(name)

    def _build_settings_ui(self):
        self._add_setting_widget("Use fast scene delegate", "/persistent/omnihydra/useFastSceneDelegate", omni.ui.CheckBox)
        self._add_setting_widget("Use scene graph instancing", "/persistent/omnihydra/useSceneGraphInstancing", omni.ui.CheckBox)
        self._add_setting_widget("Use skel adapter", "/persistent/omnihydra/useSkelAdapter", omni.ui.CheckBox)
        self._add_setting_widget("Use skel adapter blendshape", "/persistent/omnihydra/useSkelAdapterBlendShape", omni.ui.CheckBox)
        self._add_setting_widget("Use skel adapter deform graph", "/persistent/omnihydra/useSkelAdapterDeformGraph", omni.ui.CheckBox)
        self._add_setting_widget("Use cached xform time samples", "/persistent/omnihydra/useCachedXformTimeSamples", omni.ui.CheckBox)
        self._add_setting_widget("Use async RenderGraph in _PullFromRingBuffer", "/persistent/omnihydra/useAsyncRenderGraph", omni.ui.CheckBox)
        self._add_setting_widget("Use fast xform path from Fabric to renderer", "/persistent/rtx/hydra/readTransformsFromFabricInRenderDelegate", omni.ui.CheckBox)

    def _build_usdrt_settings_ui(self):
        self._add_setting_widget("Use Fabric World Bounds", "/app/omni.usd/useFabricWorldBounds", omni.ui.CheckBox)
        self._add_setting_widget("Use camera-based priority sorting", "/app/usdrt/scene_delegate/useWorldInterestBasedSorting", omni.ui.CheckBox)
        self._add_setting_widget("GPU Memory Budget %", "/app/usdrt/scene_delegate/gpuMemoryBudgetPercent", omni.ui.FloatSlider, min=0., max=100.)
        with omni.ui.CollapsableFrame("Geometry Streaming"):
            with omni.ui.VStack():
                self._add_setting_widget("Enable geometry streaming", "/app/usdrt/scene_delegate/geometryStreaming/enabled", omni.ui.CheckBox)
                self._add_setting_widget("Enable proxy cubes for unloaded prims", "/app/usdrt/scene_delegate/enableProxyCubes", omni.ui.CheckBox)
                self._add_setting_widget("Solid Angle Loading To Load in first batch", "/app/usdrt/scene_delegate/geometryStreaming/solidAngleToLoadInFirstChunk", omni.ui.FloatSlider, min=0, max=1, step=0.0001)
                self._add_setting_widget("Solid Angle Loading Limit", "/app/usdrt/scene_delegate/geometryStreaming/solidAngleLimit", omni.ui.FloatSlider, min=0, max=1, step=0.0001)
                self._add_setting_widget("Solid Angle Loading Limit Divider", "/app/usdrt/scene_delegate/geometryStreaming/solidAngleLimitDivider", omni.ui.FloatSlider, min=1, max=10000, step=100)
                self._add_setting_widget("Number of frames between batches", "/app/usdrt/scene_delegate/numFramesBetweenLoadBatches", omni.ui.IntSlider, min=1, max=50)
                self._add_setting_widget("Number of vertices to load per batch", "/app/usdrt/scene_delegate/geometryStreaming/numberOfVerticesToLoadPerChunk", omni.ui.IntSlider, min=1, max=5000000)
                self._add_setting_widget("Number of vertices to unload per batch", "/app/usdrt/scene_delegate/geometryStreaming/numberOfVerticesToUnloadPerChunk", omni.ui.IntSlider, min=1, max=5000000)
                self._add_setting_widget("Number of cubes proxy instances per instancer", "/app/usdrt/scene_delegate/proxyInstanceBucketSize", omni.ui.IntSlider, min=0, max=200000)
        with omni.ui.CollapsableFrame("Memory Budget Experimental Stability Values"):
            with omni.ui.VStack():
                self._add_setting_widget("Update Delay in microseconds", "/app/usdrt/scene_delegate/updateDelayInMicroSeconds", omni.ui.IntSlider, min=0, max=1000000)
                self._add_setting_widget("GPU Memory Deadzone %", "/app/usdrt/scene_delegate/gpuMemoryBudgetDeadZone", omni.ui.FloatSlider, min=0., max=50.)
                self._add_setting_widget("GPU Memory Filter Damping", "/app/usdrt/scene_delegate/gpuMemoryFilterDamping", omni.ui.FloatSlider, min=0., max=2.)
        with omni.ui.CollapsableFrame("Memory Limit Testing (pretend values)"):
            with omni.ui.VStack():
                self._add_setting_widget("Enable memory testing", "/app/usdrt/scene_delegate/testing/enableMemoryTesting", omni.ui.CheckBox)
                self._add_setting_widget("GPU Memory Available", "/app/usdrt/scene_delegate/testing/availableDeviceMemory", omni.ui.FloatSlider, min=0., max=float(1024*1024*1024*8))
                self._add_setting_widget("GPU Memory Total", "/app/usdrt/scene_delegate/testing/totalDeviceMemory", omni.ui.FloatSlider, min=0., max=float(1024*1024*1024*8))
        with omni.ui.CollapsableFrame("USD Population (reload scene to see effect)"):
            with omni.ui.VStack():
                self._add_setting_widget("Read Curves", "/app/usdrt/population/utils/readCurves", omni.ui.CheckBox)
                self._add_setting_widget("Read Materials", "/app/usdrt/population/utils/readMaterials", omni.ui.CheckBox)
                self._add_setting_widget("Read Lights", "/app/usdrt/population/utils/readLights", omni.ui.CheckBox)
                self._add_setting_widget("Read Primvars", "/app/usdrt/population/utils/readPrimvars", omni.ui.CheckBox)
                self._add_setting_widget("Enable Subcomponent merging", "/app/usdrt/population/utils/mergeSubcomponents", omni.ui.CheckBox)
                self._add_setting_widget("Enable Instance merging", "/app/usdrt/population/utils/mergeInstances", omni.ui.CheckBox)
                self._add_setting_widget("Enable Material merging", "/app/usdrt/population/utils/mergeMaterials", omni.ui.CheckBox)
                self._add_setting_widget("Single-threaded population", "/app/usdrt/population/utils/singleThreaded", omni.ui.CheckBox)
                self._add_setting_widget("Disable Light Scaling (non-RTX delegates)", "/app/usdrt/scene_delegate/disableLightScaling", omni.ui.CheckBox)
                self._add_setting_widget("Use Fabric Scene Graph Instancing", "/app/usdrt/population/utils/handleSceneGraphInstances", omni.ui.CheckBox)
                self._add_setting_widget("Use Hydra BlendShape", "/app/usdrt/scene_delegate/useHydraBlendShape", omni.ui.CheckBox)
                self._add_setting_widget("Number of IO threads", "/app/usdrt/population/utils/ioBoundThreadCount", omni.ui.IntSlider, min=1, max=16)

    def _build_section(self, name, build_func):
        with omni.ui.CollapsableFrame(name, height=0):
            with omni.ui.HStack():
                omni.ui.Spacer(width=20, height=5)
                with omni.ui.VStack():
                    build_func()

    def _build_ui(self):
        with omni.ui.ScrollingFrame(horizontal_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_ALWAYS_OFF,
                                    vertical_scrollbar_policy=omni.ui.ScrollBarPolicy.SCROLLBAR_AS_NEEDED):
            with omni.ui.VStack():
                self._build_section("Settings (requires stage reload)", self._build_settings_ui)
                self._build_section("USDRT Hydra Settings", self._build_usdrt_settings_ui)
