"""Imports all files of a given filetype from a selected folder
Each file gets a principal parent layer with the filename as layer name
Layer structure in the original files is preserved as sublayers
Same set of file names re-imported gets new layer set with _Copy_XX suffix.
Script by Mitch Heynick, version 30.10.18 - Windows or Mac."""

import rhinoscriptsyntax as rs
import scriptcontext as sc
import Rhino,os


def OnMac():
    return Rhino.Runtime.HostUtils.RunningOnOSX

def AddParentToLayerName(layer, parent_layer):
    #adds the parent layer name to the sub layer name to make it unique
    return "{}::{}".format(parent_layer,layer)

def GetAllFileNames(folder,ext,alt_ext):
    #finds all files of chosen filetype in the selected folder
    #returns list of tuples (filename,0) - normal ext. or (filename,1) - alt. ext.
    #returns empty list if no matching filenames found
    filenames=[]
    for filename in os.listdir(folder):
        if filename.endswith(ext):
            filenames.append((filename[:-len(ext)],0))
        elif alt_ext and filename.endswith(alt_ext):
            filenames.append((filename[:-len(alt_ext)],1))
    return filenames
    
def AddNumberSuffix(name,copy_count):
    return "{}_Copy_{}".format(name,copy_count)

def ProcessFolder(folder,ext,alt_ext,copy_count):
    
    filenames=GetAllFileNames(folder,ext,alt_ext)
    if not filenames: return
    
    obj_count=0
    file_count=0
    raise_copy_count=False
    for filename in filenames:
        if not filename[1]:
            #normal extension
            fullpath=os.path.join(folder,filename[0]+ext).lower()
        else:
            #alt extension
            fullpath=os.path.join(folder,filename[0]+alt_ext).lower()
        
        #get existing layer set
        exist_layers=set(rs.LayerNames())
        #import file
        print "Importing " + filename[0]
        rs.Command('_-Import "{}" _Enter'.format(fullpath), False)
        if rs.LastCommandResult()==0: file_count+=1
        lco = rs.LastCreatedObjects(False)
        if lco: obj_count+=len(lco)
        #get all newly imported layers
        imp_layers=set(rs.LayerNames()).difference(exist_layers)
        #Some file formats (STEP) can be "layerless"! imp_layers will have len==0
        
        #create new parent layer with filename
        if not rs.IsLayer(filename[0]):
            p_layer=rs.AddLayer(filename[0])
        else:
            #add suffix to make parent layer unique
            p_layer=rs.AddLayer(AddNumberSuffix(filename[0],copy_count))
            raise_copy_count=True
        #try to transfer layers to new parent
        if len(imp_layers)==0:
            #need to transfer imported objects to newly created parent layer
            rs.ObjectLayer(lco,p_layer)
        else:
            for imp_layer in imp_layers:
                try:
                    if not rs.ParentLayer(imp_layer):
                        rs.ParentLayer(imp_layer,p_layer)
                except: pass
                
    return True,file_count,obj_count,raise_copy_count

def BatchImportWithSublayers():
    #UI
    fTypes=["Rhino", "STL", "DXF", "DWG", "PDF", "IGES", "STEP", "SolidWorks"]
    if "BI_PrevFolder" in sc.sticky: prevFolder = sc.sticky["BI_PrevFolder"]
    else: prevFolder=rs.WorkingFolder()
    if "BI_PrevFileType" in sc.sticky: prevFileType = sc.sticky["BI_PrevFileType"]
    else: prevFileType=fTypes[0]
    #copy count data stored in document (string data!!)
    if rs.IsDocumentData():
        copy_count=rs.GetDocumentData("BI_CopyCount","Current_Count")
    else:
        copy_count=0
    if not copy_count:
        rs.SetDocumentData("BI_CopyCount","Current_Count","0")
    copy_count=rs.GetDocumentData("BI_CopyCount","Current_Count")
    
    msg="Select folder to process"
    title="Batch Import"
    folder = rs.BrowseForFolder(message=msg)
    if not folder : return
    
    msg="Select file type to import"
    fType = rs.ListBox(fTypes,msg,title,prevFileType)
    if not fType: return
    
    alt_ext=None
    if fType==fTypes[0]:
        ext = ".3dm"
    elif fType==fTypes[6]:
        ext = ".stp" ; alt_ext=".step"
    elif fType==fTypes[7]:
        ext = ".sldprt"
    elif fType==fTypes[5]:
        ext = ".igs" ; alt_ext=".iges"
    else:
        ext = "."+fType.lower()
    
    #Process the folder
    total_count=0
    rs.EnableRedraw(False)
    rc,file_count,obj_count,raise_cc=ProcessFolder(folder,ext,alt_ext,copy_count)
    rs.Redraw()
    if not rc:
        msg="No files of chosen type found in folder!"
    else:
        msg="{} files with {} total objects imported".format(file_count,obj_count)
    rs.MessageBox(msg, 0, "File Import")
    rs.UnselectAllObjects()
    
    if raise_cc:
        copy_count=str(int(copy_count)+1)
        rs.SetDocumentData("BI_CopyCount","Current_Count",copy_count)
    #store last used file type, folder
    sc.sticky["BI_PrevFolder"] = folder
    sc.sticky["BI_PrevFileType"] = fType
BatchImportWithSublayers()