bl_info = {
    "name" : "Assign Material pass_index",
    "description": "An operator for assigning pass_index values to all materials on selected objects",
    "author" : "Robert Forsman <blender@thoth.purplefrog.com>",
    "version": (0,1),
    "blender": (2,74,0),
    "location": "not on menus",
    "warning": "",
    "wiki_url": None, # "http://wiki.blender.org/index.php/Extensions:2.6/Py/Scripts/?",
    "category": "3D View",
}

import bpy

# idea by bubblebobble in IRC



def assign_material_pass_index(objects, buffer_name, rewrite_all=False):
    mats = set()
    for obj in objects:
        if hasattr(obj.data, 'materials'):
            for m in obj.data.materials:
                if m is not None:
                    mats.add(m)

    if rewrite_all:
        old_pi = set()
    else:
        old_pi = set( [ m.pass_index for m in bpy.data.materials] )
    report = "pass\tname\n"
    #report = report+"%r\n"%old_pi
    cursor=1
    for m in mats:
        while cursor in old_pi:
            cursor = cursor +1
        if rewrite_all or 0==m.pass_index:
            m.pass_index = cursor
        else:
            cursor = max(m.pass_index, cursor)
        old_pi.add(m.pass_index)
        report = report + ( "%d\t%s\n" % (m.pass_index, m.name) )

    if False:
        print(report)
    else:
        tb = bpy.data.texts.get(buffer_name)
        if tb is None:
            tb = bpy.data.texts.new(buffer_name)
        tb.from_string(report)


def thing_with_name(things, name):
    for t in things:
        if t.name==name:
            return t
    return None

def build_material_masks_for_compositor(scn):
    compositor = scn.node_tree

    file_output = compositor.nodes.get("File Output")
    if file_output is None:
        file_output = compositor.nodes.new('CompositorNodeOutputFile')
        file_output.name = "File Output"

    pass_indices = set([ m.pass_index for m in bpy.data.materials if m.pass_index !=0])

    layers_node = compositor.nodes['Render Layers']
    index_output = thing_with_name(layers_node.outputs, "IndexMA")

    print(pass_indices)
    j=0
    for i in pass_indices:
        name = "pass %d mask"%i
        mask_node = compositor.nodes.get(name)
        if mask_node is None:
            mask_node = compositor.nodes.new('CompositorNodeIDMask')
            mask_node.name = name
        mask_node.location = (100, -150*j)
        mask_node.index = i
        compositor.links.new(index_output, mask_node.inputs[0])
        if j+1>= len(file_output.inputs):
            fo_socket = file_output.file_slots.new("id %d"%i)
        else:
            fo_socket = file_output.inputs[j+1]
        compositor.links.new(mask_node.outputs[0], fo_socket)

        j=j+1



class RewriteMaterialPassIndex(bpy.types.Operator):
    bl_idname = "3dview.letterbox_center"
    bl_label = "Assign Material pass_index"
    bl_options = {'REGISTER', 'UNDO'}

    rewrite_all = bpy.props.BoolProperty(name="Rewrite All", default=False,
                                      description="are we rewriting the pass_index for materials that already have a pass_index!=0 ?")
    buffer_name = bpy.props.StringProperty(name="Name of Text Buffer", default="material report",
                                           description="name of the text buffer where we will store the report")

    def execute(self, ctx):
        #print(ctx.selected_objects)
        assign_material_pass_index(ctx.selected_objects, self.buffer_name, self.rewrite_all)
        self.report({'INFO'}, "report stored in text buffer '%s'"%self.buffer_name)
        return {'FINISHED'}

class LinkMaterialIdsToFileOutputMasks(bpy.types.Operator):
    bl_idname = "3dview.letterbox_center"
    bl_label = "Link Material Ids To File Output Masks"
    bl_options = {'REGISTER', 'UNDO'}


    def execute(self, ctx):
        #print(ctx.selected_objects)
        build_material_masks_for_compositor(ctx.scene)
        return {'FINISHED'}


def register():
    bpy.utils.register_module(__name__)

def unregister():
    bpy.utils.unregister_module(__name__)

if __name__ == "__main__":
    unregister()
    register()