import omni
from pxr import UsdGeom, Usd, Vt, Kind, Sdf

from .procedural_mesh_omnigraph_utils import (
    create_procedural_mesh_cube,
    create_procedural_mesh_sphere,
    create_procedural_mesh_capsule,
    create_procedural_mesh_cylinder,
    create_procedural_mesh_cone,
    create_procedural_mesh_torus,
    create_procedural_mesh_plane,
    create_procedural_mesh_disk,
)


class CreateProceduralMeshCommand(omni.kit.commands.Command):
    def __init__(self):
        self._create = None
        self._stage = omni.usd.get_context().get_stage()
        self._selection = omni.usd.get_context().get_selection()

    def do(self):
        xForm, mesh = self._create(self._stage)
        self._prim_path = xForm.GetPath()
        # set the new prim as the active selection
        self._selection.set_prim_path_selected(self._prim_path.pathString, True, False, True, True)
        return xForm

    def undo(self):
        prim = self._stage.GetPrimAtPath(self._prim_path)
        if prim:
            self._stage.RemovePrim(self._prim_path)


class CreateProceduralMeshCubeCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_cube


class CreateProceduralMeshSphereCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_sphere


class CreateProceduralMeshCapsuleCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_capsule


class CreateProceduralMeshCylinderCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_cylinder


class CreateProceduralMeshConeCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_cone


class CreateProceduralMeshTorusCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_torus


class CreateProceduralMeshPlaneCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_plane


class CreateProceduralMeshDiskCommand(CreateProceduralMeshCommand):
    def __init__(self):
        super().__init__()
        self._create = create_procedural_mesh_disk


omni.kit.commands.register(CreateProceduralMeshCubeCommand)
omni.kit.commands.register(CreateProceduralMeshSphereCommand)
omni.kit.commands.register(CreateProceduralMeshCapsuleCommand)
omni.kit.commands.register(CreateProceduralMeshCylinderCommand)
omni.kit.commands.register(CreateProceduralMeshConeCommand)
omni.kit.commands.register(CreateProceduralMeshTorusCommand)
omni.kit.commands.register(CreateProceduralMeshPlaneCommand)
omni.kit.commands.register(CreateProceduralMeshDiskCommand)
