import weakref
import unittest
import pathlib
import omni.kit.test
import omni.ui as ui
import carb
from pathlib import Path
from pxr import Sdf, UsdShade, UsdGeom
from omni.ui.tests.test_base import OmniUiTest
from omni.kit import ui_test
from omni.kit.test_suite.helpers import (
    open_stage,
    get_test_data_path,
    arrange_windows,
    select_prims,
    wait_stage_loading,
    )


class TestMaterialThumbContextMenu(OmniUiTest):
    async def setUp(self):
        await super().setUp()
        await arrange_windows("Stage", 200)
        await open_stage(get_test_data_path(__name__, "material_context.usda"))
        await wait_stage_loading()

    async def tearDown(self):
        await super().tearDown()
        await wait_stage_loading()

    async def test_material_thumb_context_menu_copy_1_to_1(self):
        stage = omni.usd.get_context().get_stage()

        # test copy single material from Cone_01 to Sphere_00
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())

        # select source prim
        await select_prims(["/World/Cone_01"])
        await wait_stage_loading()

        # right click and "copy"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

        # select target prim
        await select_prims(["/World/Sphere_00"])
        await ui_test.human_delay()

        # right click and "paste"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Paste", offset=ui_test.Vec2(10, 10))

        # verify binding
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")


    async def test_material_thumb_context_menu_copy_2_to_1(self):
        stage = omni.usd.get_context().get_stage()

        # test copy single material from Cone_01 to Sphere_00
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())

        # select source prim
        await select_prims(["/World/Cone_01", "/World/Cone_02"])
        await wait_stage_loading()

        # right click and "copy"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

        # select target prim
        await select_prims(["/World/Sphere_00"])
        await ui_test.human_delay()

        # right click and "paste"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Paste (2 Materials)", offset=ui_test.Vec2(10, 10))

        # verify binding
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")      


    async def test_material_thumb_context_menu_copy_4_to_2(self):
        stage = omni.usd.get_context().get_stage()

        # test copy single material from Cone_01 to Sphere_00
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_01")).ComputeBoundMaterial()
        self.assertFalse(bound_material.GetPrim().IsValid())

        # select source prim
        await select_prims(["/World/Cone_01", "/World/Cone_02"])
        await wait_stage_loading()

        # right click and "copy"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

        # select target prim
        await select_prims(["/World/Sphere_00", "/World/Sphere_01"])
        await wait_stage_loading()

        # right click and "paste"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Paste (2 Materials)", offset=ui_test.Vec2(10, 10))

        # verify binding
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_00")).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")
        
        bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath("/World/Sphere_01")).ComputeBoundMaterial()
        self.assertTrue(bound_material.GetPrim().IsValid() == True)
        self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")


    async def test_material_thumb_context_menu_copy_4211_to_2(self):
        stage = omni.usd.get_context().get_stage()

        def get_bindings(prim_paths):
            bindings = []
            for prim_path in prim_paths:
                prim = stage.GetPrimAtPath(prim_path)
                bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                bindings.append(bound_material.GetPath().pathString if bound_material else "None")

            # fixme - due to get_binding_from_prims using sets which are not ordered
            return sorted(bindings)


        for index, data in enumerate([("Paste (4 Materials)", ['/World/Looks/OmniGlass_Opacity', '/World/Looks/OmniPBR_ClearCoat', '/World/Looks/PreviewSurface', '/World/Looks/PreviewSurface']),
                                      ("Paste (2 Materials)", ['/World/Looks/PreviewSurface', '/World/Looks/PreviewSurface', 'None', 'None']),
                                      ("Paste",               ['/World/Looks/OmniPBR_ClearCoat', 'None', 'None', 'None']),
                                      ("Paste",               ['/World/Looks/OmniGlass_Opacity', 'None', 'None', 'None'])]):

            # select source prims
            await select_prims(["/World/Cone_01", "/World/Cone_02", "/World/Cube_01", "/World/Cube_02"])
            await wait_stage_loading()

            # verify copy material
            materials = get_bindings(["/World/Sphere_00", "/World/Sphere_01", "/World/Sphere_03", "/World/Sphere_04"])
            self.assertEqual(materials, ['None', 'None', 'None', 'None'])

            # right click and "copy"
            widgets = ui_test.find_all("Property//Frame/**/Button[*].identifier=='preview_drop_target'")
            await widgets[index].click(right_click=True)
            await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

            # select target prim
            await select_prims(["/World/Sphere_00", "/World/Sphere_01", "/World/Sphere_03", "/World/Sphere_04"])
            await wait_stage_loading()

            # right click and "paste"
            await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
            await ui_test.select_context_menu(data[0], offset=ui_test.Vec2(10, 10))

            materials = get_bindings(["/World/Sphere_00", "/World/Sphere_01", "/World/Sphere_03", "/World/Sphere_04"])
            self.assertEqual(materials, data[1])

            omni.kit.undo.undo()


    async def test_material_thumb_context_menu_copy_1_to_many(self):
        stage = omni.usd.get_context().get_stage()

        # test copy single material from Cone_01 to Sphere_*
        for index in range(0, 64):
            bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath(f"/World/Sphere_{index:02d}")).ComputeBoundMaterial()
            self.assertFalse(bound_material.GetPrim().IsValid())

        # select source prim
        await select_prims(["/World/Cone_01"])
        await wait_stage_loading()

        # right click and "copy"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

        # select target prims
        to_select = []
        for index in range(0, 64):
            to_select.append(f"/World/Sphere_{index:02d}")
        await select_prims(to_select)

        # right click in frame to "paste"
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Materials on selected models'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        await ui_test.select_context_menu("Paste To All", offset=ui_test.Vec2(10, 10))

        # verify binding
        for index in range(0, 64):
            bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath(f"/World/Sphere_{index:02d}")).ComputeBoundMaterial()
            self.assertTrue(bound_material.GetPrim().IsValid() == True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")


    async def test_material_thumb_context_menu_copy_2_to_many(self):
        stage = omni.usd.get_context().get_stage()

        # test copy 2 materials from Cone_01, Cone_02 to Sphere_*
        for index in range(0, 64):
            bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath(f"/World/Sphere_{index:02d}")).ComputeBoundMaterial()
            self.assertFalse(bound_material.GetPrim().IsValid())

        # select source prim
        await select_prims(["/World/Cone_01", "/World/Cone_02"])
        await wait_stage_loading()

        # right click and "copy"
        await ui_test.find("Property//Frame/**/Button[*].identifier=='preview_drop_target'").click(right_click=True)
        await ui_test.select_context_menu("Copy", offset=ui_test.Vec2(10, 10))

        # select target prims
        to_select = []
        for index in range(0, 64):
            to_select.append(f"/World/Sphere_{index:02d}")
        await select_prims(to_select)

        # right click in frame to "paste"
        widget = ui_test.find("Property//Frame/**/CollapsableFrame[*].title=='Materials on selected models'")
        await widget.click(pos=widget.position+ui_test.Vec2(widget.widget.computed_content_width/2, 10), right_click=True)
        await ui_test.select_context_menu("Paste To All", offset=ui_test.Vec2(10, 10))

        # verify binding
        for index in range(0, 64):
            bound_material, _ = UsdShade.MaterialBindingAPI(stage.GetPrimAtPath(f"/World/Sphere_{index:02d}")).ComputeBoundMaterial()
            self.assertTrue(bound_material.GetPrim().IsValid() == True)
            self.assertEqual(bound_material.GetPrim().GetPrimPath().pathString, "/World/Looks/PreviewSurface")
