import c4d
from c4d import gui

ARNOLD_RENDERER = 1029988
ARNOLD_RENDERER_COMMAND = 1039333
ARNOLD_SHADER_NETWORK = 1033991
ARNOLD_SHADER_GV = 1033990
ARNOLD_BEAUTY_PORT_ID = 537905099
ARNOLD_DRIVER = 1030141
ARNOLD_AOV = 1030369
C4DAIN_CRYPTOMATTE = 1563242911
C4DAIN_DRIVER_C4D_DISPLAY = 1927516736
C4DAIN_DRIVER_EXR = 9504161

C4DTOA_MSG_TYPE = 1000
C4DTOA_MSG_PARAM1 = 2001
C4DTOA_MSG_PARAM2 = 2002
C4DTOA_MSG_PARAM3 = 2003
C4DTOA_MSG_PARAM4 = 2004
C4DTOA_MSG_RESP1 = 2011
C4DTOA_MSG_ADD_SHADER = 1029
C4DTOA_MSG_CONNECT_ROOT_SHADER = 1033

C4DAIP_OPTIONS_AOV_SHADERS = 2113089010

CRYPTOMATTE_MATERIAL_NAME = "cryptomatte"
CRYPTOMATTE_DRIVER_NAME = "cryptomatte_driver"
DISPLAY_DRIVER_NAME = "<display driver>"

def getArnoldRenderSettings():
    rdata = doc.GetActiveRenderData()
    videopost = rdata.GetFirstVideoPost()
    while videopost:
        if videopost.GetType() == ARNOLD_RENDERER:
            return videopost;
        videopost = videopost.GetNext()
    
    if videopost is None:
        c4d.CallCommand(ARNOLD_RENDERER_COMMAND)
        
        videopost = rdata.GetFirstVideoPost()
        while videopost:
            if videopost.GetType() == ARNOLD_RENDERER:
                return videopost;
            videopost = videopost.GetNext()
            
    return None

def createArnoldShader(material, nodeId, posx, posy):
    msg = c4d.BaseContainer()
    msg.SetInt32(C4DTOA_MSG_TYPE, C4DTOA_MSG_ADD_SHADER)
    msg.SetInt32(C4DTOA_MSG_PARAM1, ARNOLD_SHADER_GV)
    msg.SetInt32(C4DTOA_MSG_PARAM2, nodeId)
    msg.SetInt32(C4DTOA_MSG_PARAM3, posx)
    msg.SetInt32(C4DTOA_MSG_PARAM4, posy)
    material.Message(c4d.MSG_BASECONTAINER, msg)
    return msg.GetLink(C4DTOA_MSG_RESP1)

def setRootShader(material, shader, rootPortId):
    msg = c4d.BaseContainer()
    msg.SetInt32(C4DTOA_MSG_TYPE, C4DTOA_MSG_CONNECT_ROOT_SHADER)
    msg.SetLink(C4DTOA_MSG_PARAM1, shader)
    msg.SetInt32(C4DTOA_MSG_PARAM2, 0)
    msg.SetInt32(C4DTOA_MSG_PARAM3, rootPortId)
    material.Message(c4d.MSG_BASECONTAINER, msg)
    return msg.GetBool(C4DTOA_MSG_RESP1)
  
def createCryptomatteShader():
    # create material
    mat = c4d.BaseMaterial(ARNOLD_SHADER_NETWORK)
    if mat is None:
        raise Exception("Failed to create material")
    mat.SetName(CRYPTOMATTE_MATERIAL_NAME)
    doc.InsertMaterial(mat)

    # create shaders
    shader = createArnoldShader(mat, C4DAIN_CRYPTOMATTE, 150, 100)
    if shader is None:
        raise Exception("Failed to create the cryptomatte shader")

    # set shader parameters
    shader.SetName("cryptomatte")
   
    # connect shaders
    setRootShader(mat, shader, ARNOLD_BEAUTY_PORT_ID)
    
    return mat;
    
def addAov(driver, aovName):
    # skip when the aov is already added
    aov = driver.GetDown()
    while aov:
        if aov.GetName() == aovName:
            return;
        aov = aov.GetNext()
        
    # create the aov
    aov = c4d.BaseObject(ARNOLD_AOV)
    aov.SetName(aovName)
    aov.InsertUnderLast(driver)    
    
def main():
    print "[C4DtoA] setup cryptomatte"

    # find the Arnold video post data    
    arnoldRenderSettings = getArnoldRenderSettings()
    if arnoldRenderSettings is None:
        raise BaseException("Failed to find Arnold render settings")

    # set Arnold as the active render engine
    rdata = doc.GetActiveRenderData()
    rdata[c4d.RDATA_RENDERENGINE] = 1029988

    # create a cryptomatte shader
    print "creating cryptomatte shader"
    material = doc.SearchMaterial(CRYPTOMATTE_MATERIAL_NAME)
    if material is None:
        material = createCryptomatteShader()
    
    # add the shader to the AOV shader list
    print "adding cryptomatte shader to the AOV shader list"
    aovShaders = arnoldRenderSettings.GetData().GetData(C4DAIP_OPTIONS_AOV_SHADERS)
    isMaterialAdded = False
    for i in range(0,aovShaders.GetObjectCount()):
        if aovShaders.ObjectFromIndex(doc, i) == material:
            isMaterialAdded = True
            break
    if isMaterialAdded is False:
        aovShaders.InsertObject(material, 0)
        arnoldRenderSettings[C4DAIP_OPTIONS_AOV_SHADERS] = aovShaders
    
    # create an EXR driver
    print "creating cryptomatte EXR driver"
    exrDriver = doc.SearchObject(CRYPTOMATTE_DRIVER_NAME)
    if exrDriver is None:
        exrDriver = c4d.BaseObject(ARNOLD_DRIVER)
        exrDriver.SetName(CRYPTOMATTE_DRIVER_NAME)
        exrDriver[c4d.C4DAI_DRIVER_TYPE] = C4DAIN_DRIVER_EXR
        exrDriver[c4d.C4DAI_DRIVER_MERGE_AOVS] = True
        doc.InsertObject(exrDriver)
          
    # add the cryptomatte AOVs to the driver
    print "adding AOVs to the cryptomatte driver"
    addAov(exrDriver, "crypto_asset")
    addAov(exrDriver, "crypto_object")
    addAov(exrDriver, "crypto_material")

    # update the scene
    c4d.EventAdd()

    print "[C4DtoA] setup cryptomatte done"

if __name__=='__main__':
    main()