"""
Extract mesh from USD/USDC and export as OBJ file.

Usage:
    python migrate/usdc_to_obj.py --input path/to/file.usdc --output path/to/output/
"""
import argparse
import os
import numpy as np
from pxr import Usd, UsdGeom


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", default=".")
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)
    obj_path = os.path.join(args.output, "Aligned_obj.obj")

    stage = Usd.Stage.Open(args.input)
    vert_offset = 0

    with open(obj_path, "w") as f:
        f.write("# Exported from USD\n\n")

        for prim in stage.Traverse():
            if prim.GetTypeName() != "Mesh":
                continue

            mesh = UsdGeom.Mesh(prim)
            points = mesh.GetPointsAttr().Get()
            if not points:
                continue

            face_counts = mesh.GetFaceVertexCountsAttr().Get()
            face_indices = mesh.GetFaceVertexIndicesAttr().Get()

            name = str(prim.GetPath()).replace("/", "_").strip("_")
            f.write(f"o {name}\n")
            print(f"  Mesh: {prim.GetPath()} ({len(points)} verts)")

            for p in points:
                f.write(f"v {p[0]:.8f} {p[1]:.8f} {p[2]:.8f}\n")

            if face_counts and face_indices:
                idx = 0
                for fc in face_counts:
                    verts = []
                    for j in range(fc):
                        verts.append(face_indices[idx] + 1 + vert_offset)
                        idx += 1
                    # Triangulate: fan from first vertex
                    for j in range(1, len(verts) - 1):
                        f.write(f"f {verts[0]} {verts[j]} {verts[j+1]}\n")

            vert_offset += len(points)

    print(f"\nSaved: {obj_path} ({vert_offset} total vertices)")


if __name__ == "__main__":
    main()