import hou
from pxr import Usd, UsdGeom, Gf, Vt, Sdf

# Access the USD stage from the LOPS context
stage = hou.pwd().editableStage()

# Function to create a box mesh with ST attributes
def create_box_mesh(stage, path, translation):
    mesh = UsdGeom.Mesh.Define(stage, path)

    # Define the vertices of a box
    points = Vt.Vec3fArray([
        (-1.0, -1.0, -1.0), (1.0, -1.0, -1.0), (1.0, 1.0, -1.0), (-1.0, 1.0, -1.0),
        (-1.0, -1.0, 1.0), (1.0, -1.0, 1.0), (1.0, 1.0, 1.0), (-1.0, 1.0, 1.0)
    ])

    # Translate points
    points = Vt.Vec3fArray([p + translation for p in points])

    # Define the face vertex counts and indices for a box (left-handed winding order)
    face_vertex_counts = Vt.IntArray([4, 4, 4, 4, 4, 4])
    face_vertex_indices = Vt.IntArray([
        0, 3, 2, 1,  # Bottom face
        4, 5, 6, 7,  # Top face
        0, 1, 5, 4,  # Front face
        2, 3, 7, 6,  # Back face
        0, 4, 7, 3,  # Left face
        1, 2, 6, 5   # Right face
    ])

    # Define the ST coordinates
    st = Vt.Vec2fArray([
        (0, 0), (1, 0), (1, 1), (0, 1),
        (0, 0), (1, 0), (1, 1), (0, 1)
    ])

    # Set the attributes on the mesh
    mesh.GetPointsAttr().Set(points)
    mesh.GetFaceVertexCountsAttr().Set(face_vertex_counts)
    mesh.GetFaceVertexIndicesAttr().Set(face_vertex_indices)

    # Create and set the ST attribute
    st_attr = mesh.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.varying)
    st_attr.Set(st)

    return mesh

# Function to get mesh data including ST attributes
def get_mesh_data(mesh_prim):
    print(f"processing prim: {mesh_prim}")
    if mesh_prim:
        points = mesh_prim.GetPointsAttr().Get()
        face_vertex_counts = mesh_prim.GetFaceVertexCountsAttr().Get()
        face_vertex_indices = mesh_prim.GetFaceVertexIndicesAttr().Get()
        primvars_api = UsdGeom.PrimvarsAPI(mesh_prim)
        st = primvars_api.GetPrimvar("st").Get()
        return points, face_vertex_counts, face_vertex_indices, st

# Function to merge multiple mesh prims into a single mesh prim
def merge_mesh_prims(stage, prim_paths, new_prim_path):
    combined_points = Vt.Vec3fArray()
    combined_face_vertex_counts = Vt.IntArray()
    combined_face_vertex_indices = Vt.IntArray()
    combined_st = Vt.Vec2fArray()

    offset = 0

    for prim_path in prim_paths:
        mesh_prim = stage.GetPrimAtPath(prim_path)
        points, face_vertex_counts, face_vertex_indices, st = get_mesh_data(UsdGeom.Mesh(mesh_prim))

        combined_points = Vt.Vec3fArray(list(combined_points) + list(points))
        combined_face_vertex_counts = Vt.IntArray(list(combined_face_vertex_counts) + list(face_vertex_counts))

        adjusted_face_vertex_indices = [index + offset for index in face_vertex_indices]
        combined_face_vertex_indices = Vt.IntArray(list(combined_face_vertex_indices) + adjusted_face_vertex_indices)

        combined_st = Vt.Vec2fArray(list(combined_st) + list(st))

        offset += len(points)

    # Create a new mesh prim
    new_mesh_prim = UsdGeom.Mesh.Define(stage, new_prim_path)

    # Set the combined points, face_vertex_counts, face_vertex_indices, and ST attributes
    new_mesh_prim.GetPointsAttr().Set(combined_points)
    new_mesh_prim.GetFaceVertexCountsAttr().Set(combined_face_vertex_counts)
    new_mesh_prim.GetFaceVertexIndicesAttr().Set(combined_face_vertex_indices)

    # Create and set the ST attribute
    primvars_api = UsdGeom.PrimvarsAPI(new_mesh_prim)
    st_attr = primvars_api.CreatePrimvar("st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying)
    st_attr.Set(combined_st)

    print(f"New merged mesh prim created at: {new_prim_path}")

# Create two box meshes for testing
#create_box_mesh(stage, "/World/Box1", Gf.Vec3f(-2, 0, 0))
#create_box_mesh(stage, "/World/Box2", Gf.Vec3f(2, 0, 0))

# List of prim paths to merge
prim_paths = ["/World/auto", "/World/single"]
new_prim_path = "/New/MergedBox"

# Merge the mesh prims
merge_mesh_prims(stage, prim_paths, new_prim_path)