__copyright__ = "Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved."
__license__ = """
NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
"""

import os
import pathlib
import platform
import sys
import time
import unittest

from . import tc_logger


def setUpModule():
    # pre-load USD before anything else to prevent crash
    # from weird carb conflict
    # FIXME ^^^
    try:
        from . import update_path_and_load_usd

        update_path_and_load_usd()
    except ImportError:
        # not needed in Kit
        pass


try:
    import omni.kit.test

    TestClass = omni.kit.test.AsyncTestCase
    TEST_DIR = str(pathlib.Path(__file__).parent.parent.parent.parent)
except ImportError:
    from . import TestRtScenegraph

    TestClass = TestRtScenegraph
    TEST_DIR = "."


class TestUsdPrim(TestClass):
    @tc_logger
    def test_has_attribute(self):
        from usdrt import Sdf, Usd, UsdGeom

        # print("Attach to {}".format(os.getpid()))
        # time.sleep(10)

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")

        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasAttribute(UsdGeom.Tokens.xformOpOrder))
        self.assertFalse(prim.HasAttribute(UsdGeom.Tokens.radius))

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)
        self.assertTrue(prim.HasAttribute(UsdGeom.Tokens.faceVertexIndices))

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.HasAttribute(UsdGeom.Tokens.radius))

    @tc_logger
    def test_get_attribute(self):
        from usdrt import Sdf, Usd, UsdGeom

        # print("Attach to {}".format(os.getpid()))
        # time.sleep(10)

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)

        attr = prim.GetAttribute(UsdGeom.Tokens.doubleSided)

        self.assertTrue(attr)
        self.assertTrue(attr.GetName() == UsdGeom.Tokens.doubleSided)

        # TODO get value

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.GetAttribute(UsdGeom.Tokens.radius))

    @tc_logger
    def test_get_attributes(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)

        attributes = prim.GetAttributes()
        self.assertTrue(len(attributes) == 12)

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.GetAttributes())

    @tc_logger
    def test_create_attribute(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.DefinePrim(Sdf.Path("/prim"))
        self.assertFalse(prim.HasAttribute("myAttr"))

        attr = prim.CreateAttribute("myAttr", Sdf.ValueTypeNames.Int, True)
        self.assertTrue(attr)
        self.assertTrue(prim.HasAttribute("myAttr"))
        self.assertTrue(attr.GetName() == "myAttr")

    @tc_logger
    def test_create_attribute_invalid(self):
        # Ensure creating attribute on invalid prim does not crash
        from usdrt import Gf, Rt, Sdf, Usd

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.IsValid())

        # this should not crash
        attr = prim.CreateAttribute("testAttr", Sdf.ValueTypeNames.Int, True)
        self.assertFalse(attr)

        # this should not crash
        xformable = Rt.Xformable(prim)
        attr = xformable.CreateWorldPositionAttr(Gf.Vec3d(0, 200, 0))
        self.assertFalse(attr)

    @tc_logger
    def test_has_relationship(self):
        from usdrt import Sdf, Usd, UsdGeom, UsdShade

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall/Green1SG"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasRelationship(UsdShade.Tokens.materialBinding))
        self.assertFalse(prim.HasRelationship(UsdGeom.Tokens.proxyPrim))

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.HasRelationship(UsdShade.Tokens.materialBinding))

    @tc_logger
    def test_get_relationship(self):
        from usdrt import Sdf, Usd, UsdShade

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall/Green1SG"))
        self.assertTrue(prim)

        rel = prim.GetRelationship(UsdShade.Tokens.materialBinding)

        targets = rel.GetTargets()
        self.assertEqual(len(targets), 1)
        self.assertEqual(targets[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG"))

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.GetRelationship(UsdShade.Tokens.materialBinding))

    @tc_logger
    def test_get_relationships(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall/Green1SG"))
        self.assertTrue(prim)

        relationships = prim.GetRelationships()
        self.assertEqual(len(relationships), 1)

        targets = relationships[0].GetTargets()
        self.assertEqual(len(targets), 1)
        self.assertEqual(targets[0], Sdf.Path("/Cornell_Box/Root/Looks/Green1SG"))

        # non-authored relationships from schema (proxyPrim) not included here
        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall"))
        relationships = prim.GetRelationships()
        self.assertEqual(len(relationships), 0)

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.GetRelationships())

    @tc_logger
    def test_create_relationship(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall/Green1SG"))
        self.assertTrue(prim)

        self.assertFalse(prim.HasRelationship("foobar"))
        rel = prim.CreateRelationship("foobar")
        self.assertTrue(prim.HasRelationship("foobar"))
        self.assertFalse(rel.HasAuthoredTargets())

    @tc_logger
    def test_create_relationship_invalid(self):
        # OM_81734 Ensure creating rel on invalid prim does not crash
        from usdrt import Sdf, Usd

        stage = Usd.Stage.CreateInMemory("test.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.IsValid())

        # this should not crash
        rel = prim.CreateRelationship("testrel")
        self.assertFalse(rel)

    @tc_logger
    def test_remove_property(self):
        from usdrt import Sdf, Usd, UsdGeom

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/Green_Wall"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasAttribute(UsdGeom.Tokens.extent))
        self.assertTrue(prim.RemoveProperty(UsdGeom.Tokens.extent))
        self.assertFalse(prim.HasAttribute(UsdGeom.Tokens.extent))

        # invalid prim should not cause crash
        prim = stage.GetPrimAtPath(Sdf.Path("/Invalid"))
        self.assertFalse(prim)
        self.assertFalse(prim.RemoveProperty(UsdGeom.Tokens.extent))

    @tc_logger
    def test_family(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        child = prim.GetChild("White_Wall_Back")
        self.assertTrue(child)
        self.assertTrue(child.GetPath() == Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))

        parent = prim.GetParent()
        self.assertTrue(parent)
        self.assertTrue(parent.GetPath() == Sdf.Path("/Cornell_Box/Root"))

        sibling = prim.GetNextSibling()
        self.assertTrue(sibling)
        self.assertTrue(sibling.GetPath() == Sdf.Path("/Cornell_Box/Root/Looks"))

        children = prim.GetChildren()
        i = 0
        for child in children:
            child_direct = prim.GetChild(child.GetName())
            self.assertEqual(child.GetName(), child_direct.GetName())
            i += 1

        self.assertEqual(i, 5)

    @tc_logger
    def test_get_type(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        test_type = prim.GetTypeName()
        self.assertEqual(test_type, "Xform")

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP/White_Wall_Back"))
        self.assertTrue(prim)

        test_type = prim.GetTypeName()
        self.assertEqual(test_type, "Mesh")

    @tc_logger
    def test_set_type(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        self.assertEqual(prim.GetTypeName(), "Xform")

        self.assertTrue(prim.SetTypeName("SkelRoot"))

        self.assertEqual(prim.GetTypeName(), "SkelRoot")

    @tc_logger
    def test_has_authored_type(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Emissive_square/Root"))
        self.assertTrue(prim)

        self.assertFalse(prim.HasAuthoredTypeName())

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasAuthoredTypeName())

    @tc_logger
    def test_clear_type(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasAuthoredTypeName())

        prim.ClearTypeName()

        self.assertFalse(prim.HasAuthoredTypeName())

    @tc_logger
    def test_repr_representation(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/Cornell_Box/Root/Cornell_Box1_LP"))
        self.assertTrue(prim)

        expected = "Prim(</Cornell_Box/Root/Cornell_Box1_LP>)"

        self.assertEqual(repr(prim), expected)

    @tc_logger
    def test_has_api(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)

        self.assertTrue(prim.HasAPI("ShapingAPI"))
        self.assertFalse(prim.HasAPI("DoesNotExist"))

    @tc_logger
    def test_apply_api(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)

        self.assertTrue(prim.ApplyAPI("Test1"))
        self.assertTrue(prim.HasAPI("Test1"))

        self.assertTrue(prim.AddAppliedSchema("Test2"))
        self.assertTrue(prim.HasAPI("Test2"))

    @tc_logger
    def test_remove_api(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)

        self.assertTrue(prim.ApplyAPI("Test1"))
        self.assertTrue(prim.HasAPI("Test1"))

        self.assertTrue(prim.AddAppliedSchema("Test2"))
        self.assertTrue(prim.HasAPI("Test2"))

        self.assertTrue(prim.RemoveAPI("Test1"))
        self.assertFalse(prim.HasAPI("Test1"))

        self.assertTrue(prim.RemoveAppliedSchema("Test2"))
        self.assertFalse(prim.HasAPI("Test2"))

    @tc_logger
    def test_get_applied_schemas(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)

        schemas = prim.GetAppliedSchemas()

        self.assertEqual(len(schemas), 5)
        self.assertTrue("LightAPI" in schemas)
        self.assertTrue("ShapingAPI" in schemas)
        self.assertTrue("CollectionAPI" in schemas)
        self.assertTrue("CollectionAPI:lightLink" in schemas)
        self.assertTrue("CollectionAPI:shadowLink" in schemas)

    @tc_logger
    def test_get_stage(self):
        from usdrt import Sdf, Usd

        stage = Usd.Stage.Open(TEST_DIR + "/data/usd/tests/cornell.usda")
        self.assertTrue(stage)

        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)

        # force stage2 out of scope
        stage2 = prim.GetStage()
        self.assertTrue(stage2)
        del stage2

        # ensure underlying shared data is unchanged
        prim = stage.GetPrimAtPath(Sdf.Path("/CylinderLight"))
        self.assertTrue(prim)
