# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#

import carb
import omni.usd

import omni.kit.commands
import omni.kit.window.content_browser as content
from omni.kit.window.popup_dialog import FormDialog, MessageDialog
import omni.ui as ui

from pxr import Sdf, Gf, UsdVol, UsdShade

from enum import Enum
import os
import re

from ..presets import ColormapPresets, ShaderPresets

class ComboListModel(ui.AbstractItemModel):
    class ComboListItem(ui.AbstractItem):
        def __init__(self, item):
            super().__init__()
            self.model = ui.SimpleStringModel(item)

    def __init__(self, item_list, default_index):
        super().__init__()
        self._default_index = default_index
        self._current_index = ui.SimpleIntModel(default_index)
        self._current_index.add_value_changed_fn(lambda a: self._item_changed(None))
        self._item_list = item_list
        self._items = []
        if item_list:
            for item in item_list:
                self._items.append(ComboListModel.ComboListItem(item))

    def get_item_children(self, item):
        return self._items

    def get_item_list(self):
        return self._item_list

    def get_item_value_model(self, item=None, column_id=-1):
        if item is None:
            return self._current_index
        return item.model

    def get_current_index(self):
        return self._current_index.get_value_as_int()

    def set_current_index(self, index):
        self._current_index.set_value(index)

    def get_current_string(self):
        return self._items[self._current_index.get_value_as_int()].model.get_value_as_string()

    def is_default(self):
        return self.get_current_index() == self._default_index


class IndeXAddVolume:
    class VolumeFormat(Enum):
        RAW, OPENVDB, PARTICLE_VOLUME = range(3)

    VOLUME_REPRESENTATIONS = {
        "NanoVDB": "vdb",
        "IndeX Sparse": "sparse_volume",
    }

    LABEL_ORDER_ROW_MAJOR = "Row-major (C-style)"
    LABEL_ORDER_COLUMN_MAJOR = "Column-major (Fortran-style)"

    LABEL_DEFAULT = "Default"
    LABEL_NONE = "None"

    def __init__(self, volume_format, ext, asset_file_path):
        self.volume_format = volume_format
        self.ext = ext
        self.asset_file_path = asset_file_path
        self.non_importer_field_names = []
        self.importer_name = None
        self.extra_importer_settings = {}
        self.file_format = None

    def _on_add(self, dlg):
        importer_settings = self.extra_importer_settings
        render_settings = {}

        colormap = None
        shader = None

        volume_usd_path = dlg.get_value("volume_usd_path")
        valid = Sdf.Path.IsValidPathString(volume_usd_path)
        if not valid:
            dlg.hide()
            def on_error_ok(msg_dlg):
                msg_dlg.hide()
                # Show the main dialog again after showing the error
                # message. Having two modal dialogs visible at the same
                # time doesn't seem to be possible.
                dlg.show()

            #FIXME: This pops up at the top-left, not the center
            MessageDialog(
                message=f"Invalid USD path: '{volume_usd_path}'\nReason: '{valid[1]}'\nPlease ensure that the specified path is valid.",
                title="Error",
                ok_handler=on_error_ok,
                disable_cancel_button=True
            ).show()
            return False

        # Pass through generic importer arguments directly
        for key, value in dlg.get_values().items():
            if key not in self.non_importer_field_names:
                if isinstance(value, bool):
                    value = str(value).lower()  # convert 'True' to 'true'
                importer_settings[key] = str(value)

        if self.volume_format == IndeXAddVolume.VolumeFormat.RAW:
            dims_model = dlg.get_field("dims").model
            sub_models = dims_model.get_item_children()
            dims = (
                dims_model.get_item_value_model(sub_models[0]).as_int,
                dims_model.get_item_value_model(sub_models[1]).as_int,
                dims_model.get_item_value_model(sub_models[2]).as_int,
            )

            for dim in dims:
                if dim <= 0:
                    dlg.hide()

                    def on_error_ok(msg_dlg):
                        msg_dlg.hide()
                        # Show the main dialog again after showing the error
                        # message. Having two modal dialogs visible at the same
                        # time doesn't seem to be possible.
                        dlg.show()

                    #FIXME: This pops up at the top-left, not the center
                    MessageDialog(
                        message=f"Invalid volume dimension specified: {dims}\nAll components must be positive.",
                        title="Error",
                        ok_handler=on_error_ok,
                        disable_cancel_button=True
                    ).show()
                    return False

            importer_settings["size"] = dims
            importer_settings["voxel_format"] = dlg.get_field("voxel_format").model.get_current_string()
            importer_settings["convert_zyx_to_xyz"] = (dlg.get_field("data_order").model.get_current_string() == self.LABEL_ORDER_ROW_MAJOR)

            importer_settings["type"] = "sparse_volume"
            importer_settings["importer"] = "nv::index::plugin::base_importer.Sparse_volume_importer_raw"

        elif self.volume_format == IndeXAddVolume.VolumeFormat.OPENVDB:
            volume_type = dlg.get_field("type").model.get_current_string()
            if volume_type in IndeXAddVolume.VOLUME_REPRESENTATIONS:
                importer_settings["type"] = IndeXAddVolume.VOLUME_REPRESENTATIONS[volume_type]

        elif self.volume_format == IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME:
            importer_settings["importer"] = self.importer_name
            importer_settings["type"] = "particle_volume"
            colormap = ColormapPresets.DEFAULT
            render_settings["renderingMode"] = dlg.get_field("renderingMode").model.get_current_string()

            fixed_point_width = dlg.get_value("fixedPointWidth")
            if (fixed_point_width > 0):
                importer_settings["fixed_radius"] = str(fixed_point_width / 2.0)

            subdivision_mode = dlg.get_field("subdivisionMode").model.get_current_string()
            if subdivision_mode != "current":
                carb.settings.get_settings().set("/rtx/index/overrideSubdivisionMode", subdivision_mode)
            subdivision_part_count = dlg.get_value("subdivisionPartCount")
            if subdivision_part_count > 0:
                carb.settings.get_settings().set("/rtx/index/overrideSubdivisionPartCount", subdivision_part_count)

        dlg.hide()

        colormap_field = dlg.get_field("colormap")
        if colormap_field and not colormap:
            colormap = colormap_field.model.get_current_string()
            if not ColormapPresets.has(colormap):
                colormap = None

        shader_field = dlg.get_field("shader")
        if shader_field and not shader:
            shader = shader_field.model.get_current_string()
            if not ShaderPresets.has(shader):
                shader = None

        omni.kit.commands.execute(
            "AddVolumeFileToStage",
            ext=self.ext,
            volume_path=volume_usd_path,
            asset_file_path=self.asset_file_path,
            volume_format=self.volume_format,
            importer_settings=importer_settings,
            render_settings=render_settings,
            colormap=colormap,
            shader=shader
        )

        return True

    def show(self):
        def on_okay(dlg: FormDialog):
            return self._on_add(dlg)

        def combo_field(name, label, items, default=0):
            return FormDialog.FieldDef(name, label, lambda **kwargs: ui.ComboBox(ComboListModel(items, default), **kwargs), 0)

        fields = []
        self.non_importer_field_names = []
        file_name = os.path.basename(self.asset_file_path)
        (file_root, file_ext) = os.path.splitext(file_name)

        # Replace characters with special meaning in SdfPath by underscores
        file_root_usd_safe = file_root
        for ch in ['.', ':', '[', ']', '-', ' ']:
            file_root_usd_safe = file_root_usd_safe.replace(ch, '_')

        default_usd_path = f"/World/volume_{file_root_usd_safe}"
        fields.append(FormDialog.FieldDef("volume_usd_path", "USD Path:  ", ui.StringField, default_usd_path))
        self.non_importer_field_names.append("volume_usd_path")

        if self.volume_format == IndeXAddVolume.VolumeFormat.RAW:
            volume_format_label = 'RAW'

            default_dims = [0, 0, 0]
            # Look for dimensions encoded in the filename
            m = re.match('.*_(\d+)_(\d+)_(\d+)(_.*|$)', file_root)
            if m:
                default_dims = [ int(m[1]), int(m[2]), int(m[3]) ]

            fields.append(FormDialog.FieldDef(
                "dims", "Dimensions:  ",
                lambda **kwargs: ui.MultiIntField(default_dims[0], default_dims[1], default_dims[2], **kwargs), 0))
            self.non_importer_field_names.append("dims")

            voxel_format_list = [ 'uint8', 'uint16', 'sint16', 'rgba8', 'float16', 'float32', 'float32_4' ]
            fields.append(combo_field("voxel_format", "Voxel Format:  ", voxel_format_list))
            self.non_importer_field_names.append("voxel_format")

            fields.append(combo_field("data_order", "Data Order:  ", [ self.LABEL_ORDER_ROW_MAJOR, self.LABEL_ORDER_COLUMN_MAJOR ]))
            self.non_importer_field_names.append("data_order")

        elif self.volume_format == IndeXAddVolume.VolumeFormat.OPENVDB:
            volume_format_label = 'OpenVDB'
            fields.append(FormDialog.FieldDef("input_file_field", "OpenVDB Data Field (optional):  ", ui.StringField, ""))
            fields.append(combo_field("type", "Volume Representation:  ", [ self.LABEL_DEFAULT ] +  list(self.VOLUME_REPRESENTATIONS.keys())))
            self.non_importer_field_names.append("type")
        elif self.volume_format == IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME:
            volume_format_label = 'particle'

            rendering_modes_list = ["rbf", "solid_box", "solid_cone", "solid_cylinder", "solid_sphere", "solid_disc"]
            fields.append(combo_field("renderingMode", "Rendering Mode:  ", rendering_modes_list, rendering_modes_list.index("solid_sphere")))
            self.non_importer_field_names.append("renderingMode")

            fields.append(FormDialog.FieldDef("fixedPointWidth", "Fixed Point Width:  ", ui.FloatField, 0.0))
            self.non_importer_field_names.append("fixedPointWidth")

            file_format = file_ext.replace('.', '').lower()
            self.file_format = file_format
            if file_format == 'las':
                self.importer_name = "nv::index::plugin::base_importer.Particle_volume_importer_las"
                fields.append(FormDialog.FieldDef("offset_point_mode", "Offset Point Mode:  ", ui.IntField, 2))
                fields.append(FormDialog.FieldDef("radius_min", "Radius Min:  ", ui.FloatField, 0.0))
                fields.append(FormDialog.FieldDef("radius_max", "Radius Max:  ", ui.FloatField, -1.0))
                fields.append(FormDialog.FieldDef("radius_scale", "Radius Scale:  ", ui.FloatField, 1.0))
                fields.append(FormDialog.FieldDef("scale_color", "Color Scale:  ", ui.FloatField, 1.0))
                fields.append(FormDialog.FieldDef("ignore_las_offset", "Ignore LAS Offset:  ", ui.CheckBox, True))
                fields.append(FormDialog.FieldDef("ignore_zero_points", "Ignore Zero Points:  ", ui.CheckBox, False))
            elif file_format == 'pcd':
                self.importer_name = "nv::index::plugin::base_importer.Particle_volume_importer_point_cloud_data"
                fields.append(FormDialog.FieldDef("is_little_endian", "Little Endian:  ", ui.CheckBox, False))
            elif file_format == 'ptx':
                self.importer_name = "nv::index::plugin::base_importer.Particle_volume_importer_ptx"
                fields.append(FormDialog.FieldDef("ignore_zero_points", "Ignore Zero Points:  ", ui.CheckBox, False))
                fields.append(FormDialog.FieldDef("color_type_uint8", "Color Type UInt8:  ", ui.CheckBox, True))
            elif file_format == 'e57':
                self.importer_name = "nv::index::plugin::e57_importer.Particle_volume_importer_e57"

                # Do not bake the transformation into the point data, but provide the transformation matrix
                self.extra_importer_settings["apply_transform"] = "false"

                fields.append(FormDialog.FieldDef("__ignore_dataset_transform", "Ignore transformation:  ", ui.CheckBox, False))
            else:
                carb.log_error(f"ERROR: Unsupported particle volume format '{file_format} for file '{self.asset_file_path}.")

            subdivision_modes_list = ["current", "grid", "kd_tree"]
            fields.append(combo_field("subdivisionMode", "Subdivision Mode:  ", subdivision_modes_list, subdivision_modes_list.index("current")))
            self.non_importer_field_names.append("subdivisionMode")
            fields.append(FormDialog.FieldDef("subdivisionPartCount", "Subdivision Part Count:  ", ui.IntField, 0))
            self.non_importer_field_names.append("subdivisionPartCount")
        else:
            return

        if self.volume_format != IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME:
            fields.append(combo_field("colormap", "Add Colormap:  ", [ self.LABEL_NONE ] + ColormapPresets.get_labels()))
            self.non_importer_field_names.append("colormap")
            fields.append(combo_field("shader", "Add XAC Shader:  ", [ self.LABEL_NONE ] + ShaderPresets.get_labels()))
            self.non_importer_field_names.append("shader")

        dlg = FormDialog(
            width=600,
            message=f"Adding {volume_format_label} volume asset '{file_name}' to the stage:",
            title=f"Add {volume_format_label.capitalize()} Volume Asset",
            ok_label="Add",
            ok_handler=on_okay,
            field_defs=fields,
            )

        dlg.show()


def is_openvdb_file(path):
    return os.path.splitext(path)[1].lower() == ".vdb"

def is_raw_file(path):
    return os.path.splitext(path)[1].lower() == ".raw"

def is_particle_volume_file(path):
    supported_particle_volume_formats = []
    file_formats = carb.settings.get_settings().get("/nvindex/state/supportedFileFormats")
    if file_formats:
        for f in file_formats:
            if f in ["las", "pcd", "ptx", "e57"]:
                supported_particle_volume_formats.append("." + f)

    return os.path.splitext(path)[1].lower() in supported_particle_volume_formats

def is_maybe_raw_file(path):
    return not (is_openvdb_file(path) or is_particle_volume_file(path))

class IndeXImporters:
    OPENVDB_FILE_HANDLER_NAME = "IndeX OpenVDB File Open Handler"
    RAW_FILE_HANDLER_NAME = "IndeX RAW File Open Handler"
    PARTICLE_VOLUME_FILE_HANDLER_NAME = "IndeX ParticleVolume File Open Handler"
    CONTEXT_MENU_NAME_ADD_OPENVDB = "Add OpenVDB Volume Asset to Stage"
    CONTEXT_MENU_NAME_ADD_RAW = "Add RAW Volume Asset to Stage"
    CONTEXT_MENU_NAME_ADD_PARTICLE_VOLUME = "Add Particle Volume Asset to Stage"

    def register_bindings(ext):
        content.get_instance().add_file_open_handler(
            IndeXImporters.OPENVDB_FILE_HANDLER_NAME,
            open_fn=lambda path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.OPENVDB, ext, path).show(),
            file_type=is_openvdb_file
        )

        content.get_instance().add_file_open_handler(
            IndeXImporters.RAW_FILE_HANDLER_NAME,
            open_fn=lambda path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.RAW, ext, path).show(),
            file_type=is_raw_file
        )

        content.get_instance().add_file_open_handler(
            IndeXImporters.PARTICLE_VOLUME_FILE_HANDLER_NAME,
            open_fn=lambda path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME, ext, path).show(),
            file_type=is_particle_volume_file
        )

        content.get_instance().add_context_menu(
            name=IndeXImporters.CONTEXT_MENU_NAME_ADD_RAW,
            glyph="menu_add.svg",
            click_fn=lambda name, path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.RAW, ext, path).show(),
            show_fn=is_maybe_raw_file
        )

        content.get_instance().add_context_menu(
            name=IndeXImporters.CONTEXT_MENU_NAME_ADD_OPENVDB,
            glyph="menu_add.svg",
            click_fn=lambda name, path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.OPENVDB, ext, path).show(),
            show_fn=is_openvdb_file
        )

        content.get_instance().add_context_menu(
            name=IndeXImporters.CONTEXT_MENU_NAME_ADD_PARTICLE_VOLUME,
            glyph="menu_add.svg",
            click_fn=lambda name, path: IndeXAddVolume(IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME, ext, path).show(),
            show_fn=is_particle_volume_file
        )

    def deregister_bindings(ext):
        content.get_instance().delete_file_open_handler(
            IndeXImporters.OPENVDB_FILE_HANDLER_NAME)
        content.get_instance().delete_file_open_handler(
            IndeXImporters.RAW_FILE_HANDLER_NAME)
        content.get_instance().delete_file_open_handler(
            IndeXImporters.PARTICLE_VOLUME_FILE_HANDLER_NAME)

        content.get_instance().delete_context_menu(
            IndeXImporters.CONTEXT_MENU_NAME_ADD_OPENVDB)
        content.get_instance().delete_context_menu(
            IndeXImporters.CONTEXT_MENU_NAME_ADD_RAW)
        content.get_instance().delete_context_menu(
            IndeXImporters.CONTEXT_MENU_NAME_ADD_PARTICLE_VOLUME)


class AddVolumeFileToStage(omni.kit.commands.Command):
    """
    Add an volume file to the stage.

    Args:
        ext: the calling extension
        volume_path: USD path of the volume in the stage
        asset_file_path: the volume assset being loaded
        volume_format: 'openvdb' or 'raw'
        importer_settings: dictionary containing setting to pass to IndeX importer
        render_settings: dictionary containing settings for rendering the volume
    """
    def __init__(self, ext, volume_path, asset_file_path, volume_format, importer_settings, render_settings, colormap, shader) -> None:
        self.ext = ext
        self.volume_path = volume_path
        self.asset_file_path = asset_file_path
        self.volume_format = volume_format
        self.importer_settings = importer_settings
        self.render_settings = render_settings
        self.colormap = colormap
        self.shader = shader

        self.stage = omni.usd.get_context().get_stage()
        self.volume_prim = None

    def do(self):
        if not self.ext:
            carb.log_error(f"ERROR: Tried to add volume {self.asset_file_path} when extension is already shut down")
            return False
        if not self.asset_file_path:
            carb.log_error(f"ERROR: Volume file {self.asset_file_path} is not defined")
            return False

        try:
            asset_name = "VolumeAsset"
            asset_path = f"{self.volume_path}/{asset_name}"
            material_path = f"{self.volume_path}/Material"
            colormap_path = f"{material_path}/Colormap"
            shader_path = f"{material_path}/VolumeShader"
            rel_name = "volume"

            # Create volume asset
            if self.volume_format == IndeXAddVolume.VolumeFormat.OPENVDB:
                volume_asset = UsdVol.OpenVDBAsset.Define(self.stage, asset_path)
                volume_asset_prim = volume_asset.GetPrim()
                field_name = self.importer_settings["input_file_field"]
                if field_name:
                    rel_name = field_name
                    volume_asset.GetFieldNameAttr().Set(field_name)
            elif self.volume_format == IndeXAddVolume.VolumeFormat.RAW:
                volume_asset_prim = self.stage.DefinePrim(asset_path, "FieldAsset")
                volume_asset = UsdVol.FieldAsset(volume_asset_prim)
            elif self.volume_format == IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME:
                volume_asset_prim = self.stage.DefinePrim(asset_path, "FieldAsset")
                volume_asset = UsdVol.FieldAsset(volume_asset_prim)

            volume_asset.CreateFilePathAttr().Set(self.asset_file_path)

            # Add generic importer settings as customData
            for key, value in self.importer_settings.items():
                # Skip keys that are handled explicitly
                if key not in ["input_file_field", "size", "type"]:
                    volume_asset_prim.SetCustomDataByKey("nvindex.importerSettings:" + key, value)

            # Create volume
            volume = UsdVol.Volume.Define(self.stage, self.volume_path)

            volume.CreateFieldRelationship(rel_name, volume_asset.GetPath())
            self.volume_prim = volume.GetPrim()

            # Mark the prim as to be handled by IndeX in compositing mode
            if carb.settings.get_settings().get("/nvindex/compositeRenderingMode"):
                self.volume_prim.CreateAttribute("nvindex:composite", Sdf.ValueTypeNames.Bool, custom=True).Set(True)
                self.volume_prim.CreateAttribute("omni:rtx:skip", Sdf.ValueTypeNames.Bool, custom=True).Set(True)

            if "size" in self.importer_settings:
                extent_attr = volume.CreateExtentAttr()
                extent_attr.Set([[0, 0, 0], self.importer_settings["size"]])

            if "type" in self.importer_settings:
                type_attr = volume.GetPrim().CreateAttribute("nvindex:type", Sdf.ValueTypeNames.Token, True)
                type_attr.Set(self.importer_settings["type"])

            # Add generic render settings as customData
            for key, value in self.render_settings.items():
                self.volume_prim.SetCustomDataByKey("nvindex.renderSettings:" + key, value)

            if self.colormap or self.shader:
                # Create and connect material and empty shader
                material = UsdShade.Material.Define(self.stage, material_path)
                UsdShade.MaterialBindingAPI(volume.GetPrim()).Bind(material)
                shader = UsdShade.Shader.Define(self.stage, shader_path)
                shader_output = shader.CreateOutput("volume", Sdf.ValueTypeNames.Token)
                material.CreateVolumeOutput("nvindex").ConnectToSource(shader_output)

                if self.volume_format == IndeXAddVolume.VolumeFormat.PARTICLE_VOLUME:
                    # Create default Phong material for particle volumes
                    shader.SetShaderId("nvindex:phong_gl")
                    shader.CreateImplementationSourceAttr().Set("id")
                    shader.CreateInput("ambient", Sdf.ValueTypeNames.Float3).Set((0.1, 0.1, 0.1))
                    shader.CreateInput("diffuse", Sdf.ValueTypeNames.Float3).Set((1.0, 1.0, 1.0))
                    shader.CreateInput("specular", Sdf.ValueTypeNames.Float3).Set((0.4, 0.4, 0.4))
                    shader.CreateInput("shininess", Sdf.ValueTypeNames.Float).Set(10.0)
                    shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(0.0) # indicates that particle color should be used

            if self.colormap:
                # Add colormap
                colormap = self.stage.DefinePrim(colormap_path, "Colormap")
                colormap_output = colormap.CreateAttribute("outputs:colormap", Sdf.ValueTypeNames.Token)
                shader.CreateInput("colormap", Sdf.ValueTypeNames.Token).ConnectToSource(colormap_output.GetPath())

                colormapSource = "rgbaPoints"
                colormap.CreateAttribute("colormapSource", Sdf.ValueTypeNames.Token).Set(colormapSource)

                (xPoints, rgbaPoints) = ColormapPresets.generate(self.colormap, colormap_source=colormapSource)
                colormap.CreateAttribute("xPoints", Sdf.ValueTypeNames.FloatArray).Set(xPoints)
                colormap.CreateAttribute("rgbaPoints", Sdf.ValueTypeNames.Float4Array).Set(rgbaPoints)

                colormap.CreateAttribute("domain", Sdf.ValueTypeNames.Float2).Set((0.0, 1.0))
            if self.shader:
                # Add shader source code
                shader.SetSourceCode(ShaderPresets.get(self.shader), "xac")

            return True
        except Exception as e:
            carb.log_warn(f"ERROR: Cannot add volume asset {self.asset_file_path}: {e}")
            return False

    def undo(self):
        if not self.ext:
            carb.log_error(f"Cannot undo adding volume {self.asset_file_path}. Extension is no longer loaded.")
            return False
        current_stage = omni.usd.get_context().get_stage()
        if current_stage != self.stage or self.volume_prim is None or not self.volume_prim.IsValid():
            carb.log_warn(f"Cannot undo adding volume {self.asset_file_path}. Stage is no longer loaded.")
            return False
        try:
            self.stage.RemovePrim(self.volume_prim.GetPath())
            self.volume_prim = None
            return True
        except Exception as e:
            carb.log_warn(f"ERROR: Cannot remove volume {self.asset_file_path}: {e}")
            return False


omni.kit.commands.register_all_commands_in_module(__name__)
