import os
import argparse
import subprocess
from multiprocessing import Pool
from generate_mosaic import generate_mosaic

def convert_one_file(root_folder, blender, f, n):
    subprocess.run([blender, "-b", "-P", "blend_to_obj.py", "--", "-f", f, "-n", str(n), "-r", root_folder], shell=False)
    
    file_list = os.path.join(root_folder, "Obj_Files", "processed_files.txt")
    
    with open(file_list, "a") as fl:
        fl.write(f + "\n")
        

def launch_conversion(root_folder, blender):
    mesh_folder = os.path.join(root_folder, "Mesh_Files")
    obj_folder = os.path.join(root_folder, "Obj_Files")
    
    uv_folder = os.path.join(obj_folder, "Artist_UVs")
    no_uv_folder = os.path.join(obj_folder, "No_UVs")
    
    cut_uv_folder = os.path.join(uv_folder, "Cut")
    uncut_uv_folder = os.path.join(uv_folder, "Uncut")
    
    cut_no_uv_folder = os.path.join(no_uv_folder, "Cut")
    uncut_no_uv_folder = os.path.join(no_uv_folder, "Uncut")
    
    img_folder = os.path.join(root_folder, "Img_Files")
    full_folder = os.path.join(obj_folder, "Full")

    file_list = os.path.join(obj_folder, "processed_files.txt")
    tags_file = os.path.join(obj_folder, "dataset_tags.csv")
    
    if not os.path.exists(obj_folder):
        os.mkdir(obj_folder)
    
    if not os.path.exists(uv_folder):
        os.mkdir(uv_folder)
        
    if not os.path.exists(no_uv_folder):
        os.mkdir(no_uv_folder)
        
    if not os.path.exists(cut_uv_folder):
        os.mkdir(cut_uv_folder)
    
    if not os.path.exists(uncut_uv_folder):
        os.mkdir(uncut_uv_folder)
        
    if not os.path.exists(cut_no_uv_folder):
        os.mkdir(cut_no_uv_folder)
    
    if not os.path.exists(uncut_no_uv_folder):
        os.mkdir(uncut_no_uv_folder)
        
    if not os.path.exists(full_folder):
        os.mkdir(full_folder)
    
    if not os.path.exists(file_list):
        f = open(file_list, "w")
        f.close()

    if not os.path.exists(tags_file):
        f = open(tags_file, "w")
        f.write("Cut,Filename,Object Number,Mesh Name,Chart Number,Disk,Closed,Manifold,Small\n")
        f.close()
        
    with open(file_list, "r") as fl:
        already_processed = fl.readlines()
        already_processed = [line.rstrip() for line in already_processed]
        
    files = [(f, os.path.getsize(os.path.join(mesh_folder, f))) for f in os.listdir(mesh_folder) if f not in already_processed]
    files.sort(key=lambda t: t[1])
    files = [f for f, s in files]
        
    with Pool(processes=8) as pool:
        results = []
        for i, f in enumerate(files):
            results.append(pool.apply_async(convert_one_file, (root_folder, blender, f, i)))
            
        for res in results:
            res.get()
            
    subprocess.run([blender, "-b", "-P", "renderings.py", "--", "-r", root_folder])
        
    generate_mosaic(img_folder, 5120, 2400, "dataset.png")
    subprocess.run(["cp", "dataset.png", os.path.join(obj_folder, "dataset.png")])




if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Launch blend converter")
    parser.add_argument("-r", "--root-folder", type=str, required=True, dest="root_folder")
    parser.add_argument("-b", "--blender", type=str, required=True, dest="blender")
    
    args = parser.parse_args()
    
    root_folder = os.path.abspath(args.root_folder)
    blender = args.blender
    
    launch_conversion(root_folder, blender)
    
    
