# batch_writer.py
# Modified 26th July 2018 to handle the application of compiled C++ RifPlugins.
# Malcolm Kesson
  
import os
  
class BatchWriter(object):
    BATCH_SCRIPT_NAME = 'batchrender__'
    BATCH_SCRIPT_PLUGINS_NAME = 'batchrender_plugins__'
    RIBS_PER_SCRIPT = 100
    #-------------------------------------------------
    # Constructor
    #-------------------------------------------------
    def __init__(self, proj_path, scene_name, vt_dir_name, ribs_dict, begin):
        self.proj_path = proj_path
        self.scene_name = scene_name
        self.ribs_dict = ribs_dict
        self.begin = begin
        # If necessary make the destination directory within the images folder.
        images_path = os.path.join(proj_path, 'images', scene_name + '_' + vt_dir_name)
        if not os.path.exists(images_path):
            os.mkdir(images_path)
        
    #-------------------------------------------------
    # makeBatchRenderScripts
    #-------------------------------------------------
    # Writes one or more batch render scripts. In the case of multiple scripts
    # each contains 100 lines of "prman -t:all PATH_TO_RIB" statements. If 
    # "immediate render mode" is ON the first script will automatically run but
    # other scripts must be run manually. 
    def makeBatchRenderScripts(self, rifPlugins=None):
        # When running a batch script by double clicking, the terminal
        # that opens does not acquire the values of the system environment 
        # variables - we resolve them here. 
        rmantree = os.getenv('RMANTREE')
        mayapath = os.getenv('MAYA_USER_DIR')
        
        rootdir = self.proj_path
        if os.name == "nt":
            rootdir = self.convertToWindows(rootdir)
            rootdir = rootdir.rstrip('\\')
        else:
            rootdir = rootdir.rstrip('/')
        prmanpath = os.path.join(rmantree, 'bin', 'prman')
        if os.name == "nt":
            prmanpath = self.convertToWindows(prmanpath)
                
        # sort the ribs
        keys = self.ribs_dict.keys()
        ribs = self.ribs_dict[keys[0]]
        
        if not rifPlugins == None:
            scriptname = BatchWriter.BATCH_SCRIPT_PLUGINS_NAME + self.scene_name
        else:
            scriptname = BatchWriter.BATCH_SCRIPT_NAME + self.scene_name
                    
        # Begin dividing into multiple scripts
        batchpath = self.getBatchScriptPath(scriptname, 1, len(ribs))
        f = open(batchpath, 'w')
        if os.name == "nt":
            f.write('\n')
        else:
            f.write('export RMANTREE=%s\n' % rmantree)
            f.write('export MAYA_USER_DIR=%s\n' % mayapath)
            if os.name == 'posix':
                f.write('unset LD_LIBRARY_PATH\n')
                
        rib_counter = 1
        num_batchscripts = 1
        batch_paths = []
        batch_paths.append(batchpath)
        for rib in ribs:
            if os.name == "nt":
                rib = self.convertToWindows(rib)
            if rifPlugins == None:
                f.write('"' + prmanpath + '" -cwd "' + rootdir + '" -t:all -progress "' + rib + '"\n')
            else:
                # Concatenate a string of the form,
                # '-rif "FULL_PATH/rifplugin.EXT" -rifargs "ARG1 ARG2" -rifend  -rif "FULL_PATH/rifplugin.EXT" -rifargs "ARG1 ARG2" -rifend'
                # Where EXT is either "so" or "dll".
                if os.name == "nt":
                    ext = '.dll'
                else:
                    ext = '.so'
                rifStr = ''
                for rif in rifPlugins:
                    rifStr += '-rif "' + rif[0] + ext + '" -rifargs "' + rif[1] + '" -rifend '
                f.write('"' + prmanpath + '" -cwd "' + rootdir + '" -t:all -progress ' + rifStr + '"' + rib + '"\n')
                
            if rib_counter % BatchWriter.RIBS_PER_SCRIPT == 0:       
                path = self.getBatchScriptPath(scriptname, rib_counter + 1,len(ribs))
                batch_paths.append(path)
                if len(path) > 0:
                    f.close()
                    #os.chmod(os.path.abspath(f.name), 0777)
                    f = open(path, 'w')
                    if os.name == "nt":
                        f.write('\n')
                    else:
                        f.write('export RMANTREE=%s\n' % rmantree)
                        f.write('export MAYA_USER_DIR=%s\n' % mayapath)
                        if os.name == 'posix':
                            f.write('unset LD_LIBRARY_PATH\n')
                    num_batchscripts += 1
            rib_counter += 1
        f.close()
        for path in batch_paths:
            os.chmod(path, 0777)
        return batch_paths
    #-------------------------------------------------
    # getBatchScriptPath
    #-------------------------------------------------    
    def getBatchScriptPath(self, scriptname, beginAt, total):    
        numRemaining = total - beginAt
        if numRemaining < 0:
            return ""        
        last = total
        if numRemaining >= BatchWriter.RIBS_PER_SCRIPT:
            last = beginAt + BatchWriter.RIBS_PER_SCRIPT - 1
        else:
            last = total
        #scriptname = BatchWriter.DEFAULT_BATCH_NAME + self.scene_name
        start = beginAt + self.begin - 1
        end = last + self.begin - 1
  
        scriptname += '_%d_%d' % (start, end)
        if os.name == "nt":
            scriptname += '.bat'
        batchpath = os.path.join(self.proj_path, scriptname)
        return batchpath
        
    # Utility_____________________________________________    
    def convertToWindows(self, linuxpath):
        pattern = re.compile(r"/")
        return pattern.sub(r'\\', linuxpath)