from pxr import Usd, Tf, Sdf, Gf, UsdGeom
import carb
from omni.kit.property.usd.prim_selection_payload import PrimSelectionPayload
import omni.kit

XFROM_OP_PREFIX = "xformOp:"
INVERSE_PREFIX = "!invert!"
INVERSE_XFORM_OP_PREFIX = "!invert!xformOp:"
RESET_XFORM_STACK = UsdGeom.XformOpTypes.resetXformStack
XFROM_OP_TYPE_NAME = [
    "translate",
    "scale",
    "rotateX",
    "rotateY",
    "rotateZ",
    "rotateXYZ",
    "rotateXZY",
    "rotateYXZ",
    "rotateYZX",
    "rotateZXY",
    "rotateZYX",
    "orient",
    "transform",
]

XFROM_OP_TYPE = [
    UsdGeom.XformOp.TypeTranslate,
    UsdGeom.XformOp.TypeScale,
    UsdGeom.XformOp.TypeRotateX,
    UsdGeom.XformOp.TypeRotateY,
    UsdGeom.XformOp.TypeRotateZ,
    UsdGeom.XformOp.TypeRotateXYZ,
    UsdGeom.XformOp.TypeRotateXZY,
    UsdGeom.XformOp.TypeRotateYXZ,
    UsdGeom.XformOp.TypeRotateYZX,
    UsdGeom.XformOp.TypeRotateZXY,
    UsdGeom.XformOp.TypeRotateZYX,
    UsdGeom.XformOp.TypeOrient,
    UsdGeom.XformOp.TypeTransform,
]

def is_inverse_op(op_name: str):
    return op_name.startswith(INVERSE_XFORM_OP_PREFIX)


def is_reset_xform_stack_op(op_name: str):
    return op_name == RESET_XFORM_STACK


def get_op_type_name(op_name: str):
    test_name = op_name
    if is_inverse_op(op_name):
        test_name = op_name.split(INVERSE_PREFIX, 1)[1]
    if test_name.startswith(XFROM_OP_PREFIX):
        names = test_name.split(":")
        if len(names) >= 2 and names[1] in XFROM_OP_TYPE_NAME:
            return names[1]
    return None


def get_op_type(op_name: str):
    op_type_name = get_op_type_name(op_name)
    return XFROM_OP_TYPE[XFROM_OP_TYPE_NAME.index(op_type_name)]

def get_op_name_suffix(op_name: str):
    test_name = op_name
    if is_inverse_op(op_name):
        test_name = op_name.split(INVERSE_PREFIX, 1)[1]
    if test_name.startswith(XFROM_OP_PREFIX):
        names = test_name.split(":", 2)
        if len(names) >= 3:
            return names[2]
    return None

def is_pivot_op(op_name: str):
    op_suffix = get_op_name_suffix(op_name)
    if op_suffix != None:
        if op_suffix == "pivot":
            return True
    return False

def is_valid_op_name(op_name: str):
    if is_reset_xform_stack_op(op_name):
        return True
    if get_op_type_name(op_name) is not None:
        return True
    return False


def get_op_attr_name(op_name: str):
    if not is_valid_op_name(op_name):
        return None
    if is_reset_xform_stack_op(op_name):
        return None
    if is_inverse_op(op_name):
        return op_name.split(INVERSE_PREFIX, 1)[1]
    return op_name


def get_inverse_op_Name(ori_op_name, desired_invert):
    if is_reset_xform_stack_op(ori_op_name):
        return ori_op_name
    if desired_invert and not is_inverse_op(ori_op_name):
        return INVERSE_PREFIX + ori_op_name
    if not desired_invert and is_inverse_op(ori_op_name):
        return ori_op_name.split(INVERSE_PREFIX, 1)[1]
    return ori_op_name


def get_op_precision(attr_type_name: Sdf.ValueTypeName):
    if attr_type_name in [Sdf.ValueTypeNames.Float3, Sdf.ValueTypeNames.Quatf, Sdf.ValueTypeNames.Float]:
        return UsdGeom.XformOp.PrecisionFloat
    if attr_type_name in [
        Sdf.ValueTypeNames.Double3,
        Sdf.ValueTypeNames.Quatd,
        Sdf.ValueTypeNames.Double,
        Sdf.ValueTypeNames.Matrix4d,
    ]:
        return UsdGeom.XformOp.PrecisionDouble
    if attr_type_name in [Sdf.ValueTypeNames.Half3, Sdf.ValueTypeNames.Quath, Sdf.ValueTypeNames.Half]:
        return UsdGeom.XformOp.PrecisionHalf
    return UsdGeom.XformOp.PrecisionDouble


def _add_trs_op(
    payload: PrimSelectionPayload
):

    settings = carb.settings.get_settings()
    # Retrieve the default precision
    default_xform_op_precision = settings.get("/persistent/app/primCreation/DefaultXformOpPrecision")
    if default_xform_op_precision is None:
        settings.set_default_string(
            "/persistent/app/primCreation/DefaultXformOpPrecision", "Double"
        )
        default_xform_op_precision = "Double"

    _precision = UsdGeom.XformOp.PrecisionDouble
    if default_xform_op_precision == "Double":
        _precision = UsdGeom.XformOp.PrecisionDouble
    elif default_xform_op_precision == "Float":
        _precision = UsdGeom.XformOp.PrecisionFloat
    elif default_xform_op_precision == "Half":
        _precision = UsdGeom.XformOp.PrecisionHalf

    # Retrieve the default rotation order
    default_rotation_order = settings.get("/persistent/app/primCreation/DefaultRotationOrder")
    if default_rotation_order is None:
        settings.set_default_string("/persistent/app/primCreation/DefaultRotationOrder", "XYZ")
        default_rotation_order = "XYZ"

    omni.kit.commands.execute("AddXformOp", payload=payload, precision=_precision, rotation_order = default_rotation_order, add_translate_op = True,
    add_rotateXYZ_op = True, add_orient_op=False, add_scale_op = True, add_transform_op = False, add_pivot_op = False)
