from pymaxwell5 import *
import os

def replace(path,from_dir,to_dir):
    p = path.replace('\\','/')
    f = from_dir.replace('\\','/')
    pnew = p.replace(f,to_dir)
    if p == pnew:
        return path,False
    return pnew,True
    

def editMaterial(mat,from_dir,to_dir):
    # Maps
    refmaps = mat.getMaps(False)
    for refmap in refmaps:
        refpathold = refmap.getPath()
        refpathnew,replaced = replace(refpathold,from_dir,to_dir)
        if replaced:
            refmap.setPath(refpathnew)
            print('REPLACED: %s -> %s' % (refpathold,refpathnew))
    nlayer,ok = mat.getNumLayers()
    for ilayer in range(nlayer):
        layer = mat.getLayer(ilayer)
        # IES
        emitter = layer.getEmitter()
        if not emitter.isNull():
            refpathold = emitter.getLobeIES()
            refpathnew,replaced = replace(refpathold,from_dir,to_dir)
            if replaced:
                if not emitter.setLobeIES(refpathnew):
                    raise('Cannot set IES: %s' % refpathnew)
                print('REPLACED: %s -> %s' % (refpathold,refpathnew))
        nbsdf,ok = layer.getNumBSDFs()
        for ibsdf in range(nbsdf):
            bsdf = layer.getBSDF(ibsdf)
            refl = bsdf.getReflectance()
            # IOR
            refpathold = refl.getComplexIor()
            refpathnew,replaced = replace(refpathold,from_dir,to_dir)
            if replaced:
                if not refl.setIOR(refpathnew):
                    raise('Cannot set IES: %s' % refpathnew)
                print('REPLACED: %s -> %s' % (refpathold,refpathnew))

def editMxm(mxmpath,from_dir,to_dir,dryrun):
    print('Editing %s' % mxmpath)
    mxs = Cmaxwell(mwcallback_cb)
    mat = mxs.readMaterial(mxmpath)
    if mat.isNull():
        raise('Cannot read %s' % mxmpath)
    editMaterial(mat,from_dir,to_dir)
    # write MXM
    if not dryrun:
        if not mat.write(mxmpath):
            raise('Cannot save %s' % mxmpath)
    print('')

def editMxs(mxspath,from_dir,to_dir,dryrun):
    print('Editing %s' % mxspath)
    # open MXS
    mxs = Cmaxwell(mwcallback_cb)
    if not mxs.readMXS(mxspath):
        raise('Cannot read %s' % mxspath)
    # fix MXS references
    ito = CmaxwellObjectIterator()
    obj = ito.first(mxs)
    while not obj.isNull():
        refpathold = obj.getReferencedScenePath()
        refpathnew,replaced = replace(refpathold,from_dir,to_dir)
        if replaced:
            obj.setReferencedScenePath(refpathnew)
            print('REPLACED: %s -> %s' % (refpathold,refpathnew))
        obj = ito.next()
    # fix materials
    itm = CmaxwellMaterialIterator()
    mat = itm.first(mxs)
    while not mat.isNull():
        refpathold,enabled = mat.getReference()
        if enabled:
            # MXM reference
            refpathnew,replaced = replace(refpathold,from_dir,to_dir)
            if replaced:
                mat.setReference(enabled,refpathnew)
                print('REPLACED: %s -> %s' % (refpathold,refpathnew))
        else:
            # Embedded material
            editMaterial(mat,from_dir,to_dir)
        mat = itm.next()
    # fix environment
    environment = mxs.getEnvironment()
    for lt in [IBL_LAYER_BACKGROUND, IBL_LAYER_REFLECTION, IBL_LAYER_REFRACTION, IBL_LAYER_ILLUMINATION]:
        refpathold,state,sphericalMapping,interpolate,intensity,uTile,vTile,uTileOffset,vTileOffset,ok = environment.getEnvironmentLayer(lt)
        refpathnew,replaced = replace(refpathold,from_dir,to_dir)
        if replaced:
            if not environment.setEnvironmentLayer(lt,refpathnew,state,sphericalMapping,interpolate,intensity,uTile,vTile,uTileOffset,vTileOffset):
                raise('Cannot set environment layer %i data' % lt)
            print('REPLACED: %s -> %s' % (refpathold,refpathnew))
    # write MXS
    if not dryrun:
        if not mxs.writeMXS(mxspath):
            raise('Cannot save %s' % mxspath)
    print('')

def edit(filepath,from_dir,to_dir,dryrun):
    if filepath.lower().endswith('.mxm'):
        editMxm(filepath,from_dir,to_dir,dryrun)
    elif filepath.lower().endswith('.mxs'):
        editMxs(filepath,from_dir,to_dir,dryrun)
    else:
        raise ValueError('Unhandled file type %s.' % filepath)

if __name__ == '__main__':

    # input dir
    input_dir = 'C:/scenes/maxwell/references/houston-sunnyvalley'
    # files to edit
    edit_this = ['.mxs','.mxm']
    # path to be replaced
    from_dir = 'C:/scenes/maxwell'
    # path to use instead
    to_dir = 'C:/scenes/paco'
    # test mode
    dryrun = True

    # get files to edit
    todo_list = []
    for root,dirs,filenames in os.walk(input_dir):
        for filename in filenames:
            if os.path.splitext(filename)[1] in edit_this:
                todo_list.append(os.path.join(root,filename))

    #edit all selected files
    for filepath in todo_list:
        edit(filepath,from_dir,to_dir,dryrun)
    print('DONE %s' % (dryrun and '-> test mode, noghing was actually saved' or ''))