# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""
Utility to convert a URDF into USD format.

Unified Robot Description Format (URDF) is an XML file format used in ROS to describe all elements of
a robot. For more information, see: http://wiki.ros.org/urdf

This script uses the URDF importer extension from Isaac Sim (``omni.isaac.urdf_importer``) to convert a
URDF asset into USD format. It is designed as a convenience script for command-line use. For more
information on the URDF importer, see the documentation for the extension:
https://docs.omniverse.nvidia.com/app_isaacsim/app_isaacsim/ext_omni_isaac_urdf.html


positional arguments:
  input               The path to the input URDF file.
  output              The path to store the USD file.

optional arguments:
  -h, --help            Show this help message and exit
  --headless            Force display off at all times. (default: False)
  --merge-joints, -m    Consolidate links that are connected by fixed joints. (default: False)
  --fix-base, -f        Fix the base to where it is imported. (default: False)
  --make-instanced, -i  Make the asset instanceable for efficient cloning. (default: False)

"""

from __future__ import annotations

"""Launch Isaac Sim Simulator first."""


import argparse
import contextlib

# omni-isaac-orbit
from omni.isaac.kit import SimulationApp

# add argparse arguments
parser = argparse.ArgumentParser(description="Utility to convert a URDF into USD format.")
parser.add_argument("input", type=str, help="The path to the input URDF file.")
parser.add_argument("output", type=str, help="The path to store the USD file.")
parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
parser.add_argument(
    "--merge-joints",
    "-m",
    action="store_true",
    default=False,
    help="Consolidate links that are connected by fixed joints.",
)
parser.add_argument(
    "--fix-base", "-f", action="store_true", default=False, help="Fix the base to where it is imported."
)
parser.add_argument(
    "--make-instanced",
    "-i",
    action="store_true",
    default=False,
    help="Make the asset instanceable for efficient cloning.",
)
args_cli = parser.parse_args()

# launch omniverse app
config = {"headless": args_cli.headless}
simulation_app = SimulationApp(config)


"""Rest everything follows."""

import os
import traceback

import carb
import omni.isaac.core.utils.stage as stage_utils
import omni.kit.app

from omni.isaac.orbit.sim.loaders import UrdfLoader, UrdfLoaderCfg
from omni.isaac.orbit.utils.assets import check_file_path
from omni.isaac.orbit.utils.dict import print_dict


def main():
    # check valid file path
    urdf_path = args_cli.input
    if not os.path.isabs(urdf_path):
        urdf_path = os.path.abspath(urdf_path)
    if not check_file_path(urdf_path):
        raise ValueError(f"Invalid file path: {urdf_path}")
    # create destination path
    dest_path = args_cli.output
    if not os.path.isabs(dest_path):
        dest_path = os.path.abspath(dest_path)

    # Create Urdf loader config
    urdf_loader_cfg = UrdfLoaderCfg(
        urdf_path=urdf_path,
        usd_dir=os.path.dirname(dest_path),
        usd_file_name=os.path.basename(dest_path),
        fix_base=args_cli.fix_base,
        merge_fixed_joints=args_cli.merge_joints,
        force_usd_conversion=True,
        make_instanceable=args_cli.make_instanced,
    )

    # Print info
    print("-" * 80)
    print("-" * 80)
    print(f"Input URDF file: {urdf_path}")
    print("URDF importer config:")
    print_dict(urdf_loader_cfg.to_dict(), nesting=0)
    print("-" * 80)
    print("-" * 80)

    # Create Urdf loader and import the file
    urdf_loader = UrdfLoader(urdf_loader_cfg)
    # print output
    print("URDF importer output:")
    print(f"Generated USD file: {urdf_loader.usd_path}")
    print("-" * 80)
    print("-" * 80)

    # Simulate scene (if not headless)
    if not args_cli.headless:
        # Open the stage with USD
        stage_utils.open_stage(urdf_loader.usd_path)
        # Reinitialize the simulation
        app = omni.kit.app.get_app_interface()
        # Run simulation
        with contextlib.suppress(KeyboardInterrupt):
            while True:
                # perform step
                app.update()


if __name__ == "__main__":
    try:
        # run the main execution
        main()
    except Exception as err:
        carb.log_error(err)
        carb.log_error(traceback.format_exc())
        raise
    finally:
        # close sim app
        simulation_app.close()
