import omni.usd
import omni.kit.actions.core

from pxr import Usd, UsdGeom, UsdLux, OmniAudioSchema


def get_action_name(name):
    return name.lower().replace('-', '_').replace(' ', '_')


def get_geometry_standard_prim_list(usd_context=None):
    if not usd_context:
        usd_context = omni.usd.get_context()

    stage = usd_context.get_stage()
    if stage:
        meters_per_unit = UsdGeom.GetStageMetersPerUnit(stage)
    else:
        meters_per_unit = 0.01

    geom_base = 0.5 / meters_per_unit
    geom_base_double = geom_base * 2
    geom_base_half = geom_base / 2

    all_shapes = sorted([
        (
            "Cube",
            {
                UsdGeom.Tokens.size: geom_base_double,
                UsdGeom.Tokens.extent: [(-geom_base, -geom_base, -geom_base), (geom_base, geom_base, geom_base)],
            },
        ),
        (
            "Sphere",
            {
                UsdGeom.Tokens.radius: geom_base,
                UsdGeom.Tokens.extent: [(-geom_base, -geom_base, -geom_base), (geom_base, geom_base, geom_base)],
            },
        ),
        (
            "Cylinder",
            {
                UsdGeom.Tokens.radius: geom_base,
                UsdGeom.Tokens.height: geom_base_double,
                UsdGeom.Tokens.extent: [(-geom_base, -geom_base, -geom_base), (geom_base, geom_base, geom_base)],
            },
        ),
        (
            "Capsule",
            {
                UsdGeom.Tokens.radius: geom_base_half,
                UsdGeom.Tokens.height: geom_base,
                # Create extent with command dynamically since it's different in different up-axis.
            },
        ),
        (
            "Cone",
            {
                UsdGeom.Tokens.radius: geom_base,
                UsdGeom.Tokens.height: geom_base_double,
                UsdGeom.Tokens.extent: [(-geom_base, -geom_base, -geom_base), (geom_base, geom_base, geom_base)],
            },
        ),
    ])

    shape_attrs = {}
    for name, attrs in all_shapes:
        shape_attrs[name] = attrs
    
    return shape_attrs


def get_audio_prim_list():
    return [
        ("Spatial Sound", "OmniSound", {}),
        (
            "Non-Spatial Sound",
            "OmniSound",
            {OmniAudioSchema.Tokens.auralMode: OmniAudioSchema.Tokens.nonSpatial}
        ),
        ("Listener", "OmniListener", {})
    ]


def get_light_prim_list(usd_context=None):
    if not usd_context:
        usd_context = omni.usd.get_context()

    stage = usd_context.get_stage()
    if stage:
        meters_per_unit = UsdGeom.GetStageMetersPerUnit(stage)
    else:
        meters_per_unit = 0.01

    geom_base = 0.5 / meters_per_unit
    geom_base_double = geom_base * 2

    # https://github.com/PixarAnimationStudios/USD/commit/b5d3809c943950cd3ff6be0467858a3297df0bb7
    if hasattr(UsdLux.Tokens, 'inputsIntensity'):
        return sorted([
            ("Distant Light", "DistantLight", {UsdLux.Tokens.inputsAngle: 1.0, UsdLux.Tokens.inputsIntensity: 3000}),
            ("Sphere Light", "SphereLight", {UsdLux.Tokens.inputsRadius: geom_base, UsdLux.Tokens.inputsIntensity: 30000}),
            (
                "Rect Light",
                "RectLight",
                {
                    UsdLux.Tokens.inputsWidth: geom_base_double,
                    UsdLux.Tokens.inputsHeight: geom_base_double,
                    UsdLux.Tokens.inputsIntensity: 15000,
                },
            ),
            ("Disk Light", "DiskLight", {UsdLux.Tokens.inputsRadius: geom_base, UsdLux.Tokens.inputsIntensity: 60000}),
            (
                "Cylinder Light",
                "CylinderLight",
                {UsdLux.Tokens.inputsLength: geom_base_double, UsdLux.Tokens.inputsRadius: 5, UsdLux.Tokens.inputsIntensity: 30000},
            ),
            (
                "Dome Light",
                "DomeLight",
                {UsdLux.Tokens.inputsIntensity: 1000, UsdLux.Tokens.inputsTextureFormat: UsdLux.Tokens.latlong},
            ),
        ])

    return sorted([
        ("Distant Light", "DistantLight", {UsdLux.Tokens.angle: 1.0, UsdLux.Tokens.intensity: 3000}),
        ("Sphere Light", "SphereLight", {UsdLux.Tokens.radius: geom_base, UsdLux.Tokens.intensity: 30000}),
        (
            "Rect Light",
            "RectLight",
            {
                UsdLux.Tokens.width: geom_base_double,
                UsdLux.Tokens.height: geom_base_double,
                UsdLux.Tokens.intensity: 15000,
            },
        ),
        ("Disk Light", "DiskLight", {UsdLux.Tokens.radius: geom_base, UsdLux.Tokens.intensity: 60000}),
        (
            "Cylinder Light",
            "CylinderLight",
            {UsdLux.Tokens.length: geom_base_double, UsdLux.Tokens.radius: 5, UsdLux.Tokens.intensity: 30000},
        ),
        (
            "Dome Light",
            "DomeLight",
            {UsdLux.Tokens.intensity: 1000, UsdLux.Tokens.textureFormat: UsdLux.Tokens.latlong},
        ),
    ])


def register_actions(extension_id, cls, get_self_fn):
    actions_tag = "Create Menu Actions"

    # actions
    omni.kit.actions.core.get_action_registry().register_action(
        extension_id,
        "create_prim",
        cls.on_create_prim,
        display_name="Create->Create Prim",
        description="Create Prim",
        tag=actions_tag
    )

    omni.kit.actions.core.get_action_registry().register_action(
        extension_id,
        "create_prim_camera",
        lambda: cls.on_create_prim("Camera", {}),
        display_name="Create->Create Camera Prim",
        description="Create Camera Prim",
        tag=actions_tag
    )

    omni.kit.actions.core.get_action_registry().register_action(
        extension_id,
        "create_prim_scope",
        lambda: cls.on_create_prim("Scope", {}),
        display_name="Create->Create Scope",
        description="Create Scope",
        tag=actions_tag
    )

    omni.kit.actions.core.get_action_registry().register_action(
        extension_id,
        "create_prim_xform",
        lambda: cls.on_create_prim("Xform", {}),
        display_name="Create->Create Xform",
        description="Create Xform",
        tag=actions_tag
    )

    def on_create_shape(shape_name):
        shape_attrs = get_geometry_standard_prim_list()
        cls.on_create_prim(shape_name, shape_attrs[shape_name])

    for prim in get_geometry_standard_prim_list().keys():
        omni.kit.actions.core.get_action_registry().register_action(
            extension_id,
            f"create_prim_{prim.lower()}",
            lambda p=prim: on_create_shape(p),
            display_name=f"Create->Create Prim {prim}",
            description=f"Create Prim {prim}",
            tag=actions_tag
        )

    for name, prim, attrs in get_light_prim_list():
        omni.kit.actions.core.get_action_registry().register_action(
            extension_id,
            f"create_prim_{get_action_name(name)}",
            lambda p=prim, a=attrs: cls.on_create_light(p, a),
            display_name=f"Create->Create Prim {name}",
            description=f"Create Prim {name}",
            tag=actions_tag
        )

    for name, prim, attrs in get_audio_prim_list():
        omni.kit.actions.core.get_action_registry().register_action(
            extension_id,
            f"create_prim_{get_action_name(name)}",
            lambda p=prim, a=attrs: cls.on_create_prim(p, a),
            display_name=f"Create->Create Prim {name}",
            description=f"Create Prim {name}",
            tag=actions_tag
        )

    omni.kit.actions.core.get_action_registry().register_action(
        extension_id,
        "high_quality_option_toggle",
        get_self_fn()._high_quality_option_toggle,
        display_name="Create->High Quality Option Toggle",
        description="Create High Quality Option Toggle",
        tag=actions_tag
)


def deregister_actions(extension_id):
    action_registry = omni.kit.actions.core.get_action_registry()
    action_registry.deregister_all_actions_for_extension(extension_id)
