"""
A set of utilities
"""
import omni.kit.app
import omni.usd

from pxr import Sdf, UsdGeom
import omni.graph.core as og


def get_parent(stage, parentPath):
    if parentPath:
        parentPrim = stage.GetPrimAtPath(parentPath)
        if parentPrim:
            return parentPrim, False

    # fallback to default behavior (creating a parent xform prim)
    omni.kit.commands.execute(
        "CreatePrimWithDefaultXform",
        prim_type="Xform",
        prim_path=omni.usd.get_stage_next_free_path(stage, "/" + "ProceduralMeshXform", True),
        select_new_prim=True,
        create_default_xform=False,
    )
    paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
    for path in paths:
        parentPrim = stage.GetPrimAtPath(path)
    return parentPrim, True


def create_mesh_computer(stage, prim_name: str, evaluator_name: str, node_type: str, parentPath=None):
    parentPrim, isParentOwned = get_parent(stage, parentPath)

    omni.kit.commands.execute(
        "CreatePrimWithDefaultXform",
        prim_type="Mesh",
        prim_path=omni.usd.get_stage_next_free_path(
            stage, str(parentPrim.GetPath().AppendElementString(prim_name)), False
        ),
        select_new_prim=True,
        create_default_xform=False,
    )
    paths = omni.usd.get_context().get_selection().get_selected_prim_paths()
    for path in paths:
        mesh = stage.GetPrimAtPath(path)

    helper = og.OmniGraphHelper()
    computer = helper.create_node(str(mesh.GetPath().AppendElementString(evaluator_name)), node_type)

    if isParentOwned:
        return parentPrim, mesh, computer
    else:
        return mesh, mesh, computer


def init_common_attributes(prim):
    # parameter check
    # sschirm
    parameterCheckAttr = prim.CreateAttribute("proceduralMesh:parameterCheck", Sdf.ValueTypeNames.Int)
    parameterCheckAttr.Set(0)

    # basic attributes
    pointsAttr = prim.GetAttribute("points")
    faceVertexCountsAttr = prim.GetAttribute("faceVertexCounts")
    faceVertexIndicesAttr = prim.GetAttribute("faceVertexIndices")

    # extent
    extentAttr = prim.GetAttribute("extent")


def connect_common_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect parameterCheck attribute as in and output
    helper.connect(ognPrim, "proceduralMesh:parameterCheck", computer, "inputs:parameterCheck")
    helper.connect(computer, "outputs:parameterCheck", ognPrim, "proceduralMesh:parameterCheck")

    # connect basic attributes
    helper.connect(computer, "outputs:points", ognPrim, "points")
    helper.connect(computer, "outputs:faceVertexCounts", ognPrim, "faceVertexCounts")
    helper.connect(computer, "outputs:faceVertexIndices", ognPrim, "faceVertexIndices")

    # connect to extent
    helper.connect(computer, "outputs:extent", ognPrim, "extent")


def init_crease_attributes(prim):
    creaseIndicesAttr = prim.GetAttribute("creaseIndices")
    creaseLengthsAttr = prim.GetAttribute("creaseLengths")
    creaseSharpnessesAttr = prim.GetAttribute("creaseSharpnesses")


def connect_crease_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect smooth group related attributes
    helper.connect(computer, "outputs:creaseIndices", ognPrim, "creaseIndices")
    helper.connect(computer, "outputs:creaseLengths", ognPrim, "creaseLengths")
    helper.connect(computer, "outputs:creaseSharpnesses", ognPrim, "creaseSharpnesses")


def init_corner_attributes(prim):
    cornerIndicesAttr = prim.GetAttribute("cornerIndices")
    cornerSharpnessesAttr = prim.GetAttribute("cornerSharpnesses")


def connect_corner_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect smooth group related attributes
    helper.connect(computer, "outputs:cornerIndices", ognPrim, "cornerIndices")
    helper.connect(computer, "outputs:cornerSharpnesses", ognPrim, "cornerSharpnesses")


def init_normals_facevarying_attributes(prim):
    mesh = UsdGeom.Mesh(prim)

    # authored normals only works when subdivision is none
    mesh.CreateSubdivisionSchemeAttr("none")
    normalsVar = mesh.CreatePrimvar("normals", Sdf.ValueTypeNames.Normal3fArray)
    normalsAttr = prim.GetAttribute("primvars:normals")
    normalIndicesAttr = prim.CreateAttribute("primvars:normals:indices", Sdf.ValueTypeNames.IntArray, False)

    normalsVar.SetInterpolation("faceVarying")


def connect_normals_facevarying_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect normal related attributes
    helper.connect(computer, "outputs:normals", ognPrim, "primvars:normals")
    helper.connect(computer, "outputs:normalIndices", ognPrim, "primvars:normals:indices")


def init_texture_coords_facevarying_attributes(prim):
    mesh = UsdGeom.Mesh(prim)

    # st related attributes
    stsVar = mesh.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying)
    stsAttr = prim.GetAttribute("primvars:st")
    stIndicesAttr = prim.CreateAttribute("primvars:st:indices", Sdf.ValueTypeNames.IntArray, False)


def connect_texture_coords_facevarying_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect st related attributes
    helper.connect(computer, "outputs:sts", ognPrim, "primvars:st")
    helper.connect(computer, "outputs:stIndices", ognPrim, "primvars:st:indices")


def init_texture_coords_attributes(prim):
    mesh = UsdGeom.Mesh(prim)

    # st related attributes
    stsVar = mesh.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.vertex)
    stsAttr = prim.GetAttribute("primvars:st")


def connect_texture_coords_attributes(prim, computer):
    helper = og.OmniGraphHelper()
    ognPrim = helper.attach_to_prim(prim)

    # connect st related attributes
    helper.connect(computer, "outputs:sts", ognPrim, "primvars:st")


def create_procedural_mesh_cube(
    stage, parentPath=None, xLength=100, yLength=100, zLength=100, xResolution=10, yResolution=10, zResolution=10
):
    root, cube, computer = create_mesh_computer(
        stage, "ProceduralMeshCube", "CubeEvaluator", "omni.kit.procedural.mesh.CreateMeshCube", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:xLength", xLength],
            ["inputs:yLength", yLength],
            ["inputs:zLength", zLength],
            ["inputs:xResolution", xResolution],
            ["inputs:yResolution", xResolution],
            ["inputs:zResolution", zResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(cube)
    init_crease_attributes(cube)
    init_normals_facevarying_attributes(cube)
    init_texture_coords_facevarying_attributes(cube)
    connect_common_attributes(cube, computer)
    connect_crease_attributes(cube, computer)
    connect_normals_facevarying_attributes(cube, computer)
    connect_texture_coords_facevarying_attributes(cube, computer)

    return root, cube


def create_procedural_mesh_sphere(stage, parentPath=None, diameter=100, longitudeResolution=10, latitudeResolution=10):
    root, sphere, computer = create_mesh_computer(
        stage, "ProceduralMeshSphere", "SphereEvaluator", "omni.kit.procedural.mesh.CreateMeshSphere", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:longitudeResolution", longitudeResolution],
            ["inputs:latitudeResolution", latitudeResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(sphere)
    init_texture_coords_facevarying_attributes(sphere)
    connect_common_attributes(sphere, computer)
    connect_texture_coords_facevarying_attributes(sphere, computer)

    return root, sphere


def create_procedural_mesh_capsule(
    stage,
    parentPath=None,
    diameter=100,
    yLength=200,
    longitudeResolution=10,
    hemisphereLatitudeResolution=10,
    middleLatitudeResolution=10,
):
    root, capsule, computer = create_mesh_computer(
        stage, "ProceduralMeshCapsule", "CapsuleEvaluator", "omni.kit.procedural.mesh.CreateMeshCapsule", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:yLength", yLength],
            ["inputs:longitudeResolution", longitudeResolution],
            ["inputs:hemisphereLatitudeResolution", hemisphereLatitudeResolution],
            ["inputs:middleLatitudeResolution", middleLatitudeResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(capsule)
    init_texture_coords_facevarying_attributes(capsule)
    connect_common_attributes(capsule, computer)
    connect_texture_coords_facevarying_attributes(capsule, computer)

    return root, capsule


def create_procedural_mesh_cylinder(
    stage, parentPath=None, diameter=100, height=100, longitudeResolution=10, latitudeResolution=10, capResolution=10
):
    root, cylinder, computer = create_mesh_computer(
        stage, "ProceduralMeshCylinder", "CylinderEvaluator", "omni.kit.procedural.mesh.CreateMeshCylinder", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:height", height],
            ["inputs:longitudeResolution", longitudeResolution],
            ["inputs:latitudeResolution", latitudeResolution],
            ["inputs:capResolution", capResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(cylinder)
    init_crease_attributes(cylinder)
    init_normals_facevarying_attributes(cylinder)
    init_texture_coords_facevarying_attributes(cylinder)
    connect_common_attributes(cylinder, computer)
    connect_crease_attributes(cylinder, computer)
    connect_normals_facevarying_attributes(cylinder, computer)
    connect_texture_coords_facevarying_attributes(cylinder, computer)

    return root, cylinder


def create_procedural_mesh_cone(
    stage, parentPath=None, diameter=100, height=100, longitudeResolution=10, latitudeResolution=10, capResolution=10
):
    root, cone, computer = create_mesh_computer(
        stage, "ProceduralMeshCone", "ConeEvaluator", "omni.kit.procedural.mesh.CreateMeshCone", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:height", height],
            ["inputs:longitudeResolution", longitudeResolution],
            ["inputs:latitudeResolution", latitudeResolution],
            ["inputs:capResolution", capResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(cone)
    init_crease_attributes(cone)
    init_corner_attributes(cone)
    init_normals_facevarying_attributes(cone)
    init_texture_coords_facevarying_attributes(cone)
    connect_common_attributes(cone, computer)
    connect_crease_attributes(cone, computer)
    connect_corner_attributes(cone, computer)
    connect_normals_facevarying_attributes(cone, computer)
    connect_texture_coords_facevarying_attributes(cone, computer)

    return root, cone


def create_procedural_mesh_torus(
    stage, parentPath=None, diameter=200, sectionDiameter=50, sectionTwist=0, loopResolution=10, sectionResolution=10
):
    root, torus, computer = create_mesh_computer(
        stage, "ProceduralMeshTorus", "TorusEvaluator", "omni.kit.procedural.mesh.CreateMeshTorus", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:sectionDiameter", sectionDiameter],
            ["inputs:sectionTwist", sectionTwist],
            ["inputs:loopResolution", loopResolution],
            ["inputs:sectionResolution", sectionResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(torus)
    init_texture_coords_facevarying_attributes(torus)
    connect_common_attributes(torus, computer)
    connect_texture_coords_facevarying_attributes(torus, computer)

    return root, torus


def create_procedural_mesh_disk(stage, parentPath=None, diameter=100, longitudeResolution=10, capResolution=10):
    root, disk, computer = create_mesh_computer(
        stage, "ProceduralMeshDisk", "DiskEvaluator", "omni.kit.procedural.mesh.CreateMeshDisk", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:diameter", diameter],
            ["inputs:longitudeResolution", longitudeResolution],
            ["inputs:capResolution", capResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(disk)
    init_texture_coords_attributes(disk)
    connect_common_attributes(disk, computer)
    connect_texture_coords_attributes(disk, computer)

    return root, disk


def create_procedural_mesh_plane(stage, parentPath=None, xLength=100, zLength=100, xResolution=10, zResolution=10):
    root, plane, computer = create_mesh_computer(
        stage, "ProceduralMeshPlane", "PlaneEvaluator", "omni.kit.procedural.mesh.CreateMeshPlane", parentPath
    )

    helper = og.OmniGraphHelper()
    # omnigraph node for the mesh init
    helper.set_values(
        computer,
        [
            ["inputs:xLength", xLength],
            ["inputs:zLength", zLength],
            ["inputs:xResolution", xResolution],
            ["inputs:zResolution", zResolution],
        ],
    )

    # setup mesh attributes and node inputs/outputs
    init_common_attributes(plane)
    init_texture_coords_attributes(plane)
    connect_common_attributes(plane, computer)
    connect_texture_coords_attributes(plane, computer)

    return root, plane
