import carb
from pxr import Usd, Sdf, UsdShade


class Constant:
    def __setattr__(self, name, value):
        raise Exception(f"Can't change Constant.{name}")

    ICON_SIZE = 96
    BOUND_LABEL_WIDTH = 50
    FONT_SIZE = 14.0
    SDF_PATH_INVALID = "$NONE$"
    PERSISTENT_SETTINGS_PREFIX = "/persistent"
    MIXED = "Mixed"
    MIXED_COLOR = 0xFFCC9E61
    # verify SDF_PATH_INVALID is invalid
    if Sdf.Path.IsValidPathString(SDF_PATH_INVALID):
        raise Exception(f"SDF_PATH_INVALID is Sdf.Path.IsValidPathString - FIXME")


def _populate_data(stage, material_data, collection_or_prim, material=None, relationship=None):
    def update_strength(value, strength):
        if value is None:
            value = strength
        elif value != strength:
            value = Constant.MIXED
        return value
    
    if isinstance(collection_or_prim, Usd.CollectionAPI):
        prim_path = collection_or_prim.GetCollectionPath()
    else:
        prim_path = collection_or_prim.GetPath()

    inherited = False
    strength_default = carb.settings.get_settings().get(
        Constant.PERSISTENT_SETTINGS_PREFIX + "/app/stage/materialStrength"
    )
    strength = strength_default
    material_name = Constant.SDF_PATH_INVALID

    if material:
        relationship_source_path = None
        if relationship:
            relationship_source_path = relationship.GetPrim().GetPath()
            if isinstance(collection_or_prim, Usd.CollectionAPI):
                relationship_source_path = relationship.GetTargets()[0]

        if (
            material
            and relationship
            and prim_path.pathString != relationship_source_path.pathString
        ):
            #If we have a material, but it's not assigned directly to us, it must be inherited
            inherited = True
        if relationship:
            strength = UsdShade.MaterialBindingAPI.GetMaterialBindingStrength(relationship)
        material_name = material.GetPrim().GetPath().pathString

    if not material_name in material_data["material"]:
        material_data["material"][material_name] = {
            "bound": set(),
            "inherited": set(),
            "bound_info": set(),
            "inherited_info": set(),
            "relationship": set(),
            "strength": None,
        }

    if inherited:
        material_data["material"][material_name]["inherited"].add(collection_or_prim)
        material_data["material"][material_name]["inherited_info"].add(
            (prim_path.pathString, material_name, strength)
        )
        material_data["inherited"].add(collection_or_prim)
        material_data["inherited_info"].add((prim_path.pathString, material_name, strength))
    else:
        material_data["material"][material_name]["bound"].add(collection_or_prim)
        material_data["material"][material_name]["bound_info"].add(
            (prim_path.pathString, material_name, strength)
        )
        material_data["bound"].add(collection_or_prim)
        material_data["bound_info"].add((prim_path.pathString, material_name, strength))

    if relationship:
        material_data["relationship"].add(relationship)
        material_data["material"][material_name]["relationship"].add(relationship)

    if strength is not None:
        material_data["material"][material_name]["strength"] = update_strength(
            material_data["material"][material_name]["strength"], strength
        )
        material_data["strength"] = update_strength(material_data["strength"], strength)
            

def get_binding_from_prims(stage, prim_paths):
    material_data = {"material": {}, "bound": set(), "bound_info": set(), "inherited": set(), "inherited_info": set(), "relationship": set(), "strength": None}
      
    for prim_path in prim_paths:
        if prim_path.IsPrimPath():
            prim = stage.GetPrimAtPath(prim_path)
            if prim:
                material, relationship = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                _populate_data(stage, material_data, prim, material, relationship)
                  
        elif Usd.CollectionAPI.IsCollectionAPIPath(prim_path):
            real_prim_path = prim_path.GetPrimPath()
            prim = stage.GetPrimAtPath(real_prim_path)
            binding_api = UsdShade.MaterialBindingAPI(prim)
            all_bindings = binding_api.GetCollectionBindings()
            collection_name = prim_path.pathString[prim_path.pathString.find(".collection:")+12:]
                
            #TODO: test this when we have multiple collections on a prim
            if all_bindings:
                for b in all_bindings:
                    collection = b.GetCollection()
                    if collection_name==collection.GetName():
                        relationship = b.GetBindingRel()
                        material = b.GetMaterial()
                        _populate_data(stage, material_data, collection, material, relationship)
            else:
                #If there are no bindings we want to set up defaults anyway so the widget appears with "None" assigned
                collection = Usd.CollectionAPI.Get(stage, prim_path)
                _populate_data(stage, material_data, collection)

    return material_data
