import argparse
from OCC.Core.Graphic3d import *
from OCC.Display.OCCViewer import Viewer3d
from OCC.Extend.DataExchange import read_step_file
from OCC.Extend.TopologyUtils import TopologyExplorer
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB, Quantity_NOC_WHITE
from OCC.Core.V3d import V3d_DirectionalLight
from OCC.Core.gp import gp_Dir
from glob import glob
import pathlib
from tqdm import tqdm


def render(shape, filename, width=1024, height=768, face_color_rgb=(0.2, 0.2, 0.2), edge_color_rgb=(0, 0, 0), show_face_boundary=True):
    viewer = Viewer3d()
    viewer.Create(phong_shading=True, create_default_lights=True)
    viewer.set_bg_gradient_color([255, 255, 255], [255, 255, 255])
    viewer.SetModeShaded()
    viewer.hide_triedron()
    viewer.EnableAntiAliasing()
    dir_light = V3d_DirectionalLight(gp_Dir(0, 0.5, -1), Quantity_Color(Quantity_NOC_WHITE))
    dir_light.SetEnabled(True)
    dir_light.SetIntensity(500.0)
    viewer.Viewer.AddLight(dir_light)
    viewer.Viewer.SetLightOn()

    viewer.default_drawer.EnableDrawHiddenLine()
    viewer.default_drawer.SetFaceBoundaryDraw(show_face_boundary)
    ais_context = viewer.GetContext()
    dc = ais_context.DeviationCoefficient()
    da = ais_context.DeviationAngle()
    factor = 10
    ais_context.SetDeviationCoefficient(dc / factor)
    ais_context.SetDeviationAngle(da / factor)
    topexp = TopologyExplorer(shape)
    for face in topexp.faces():
        if face is not None:
            viewer.DisplayShape(face, color=Quantity_Color(*face_color_rgb, Quantity_TOC_RGB))
    for edge in topexp.edges():
        if edge is not None:
            viewer.DisplayShape(edge, color=Quantity_Color(*edge_color_rgb, Quantity_TOC_RGB))
    viewer.FitAll()
    viewer.SetSize(width, height)
    viewer.View.Dump(str(filename))


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--input_dir", type=str, required=True, help="Input folder of STP/STEP files")
    p.add_argument("--output_dir", type=str, required=True, help="Output folder of PNG files")
    p.add_argument("--width", type=int, default=1024, help="Width of image")
    p.add_argument("--height", type=int, default=768, help="Height of image")

    args = p.parse_args()
    
    files = []
    cad_folders = sorted(glob(args.input_dir+'/*/'))
    for folder in cad_folders:
        input_path = pathlib.Path(folder)
        files += list(input_path.glob("*.st*p"))
    print(len(files))
    # files = files[36000:] # debug only (* remove *)
    output_path = pathlib.Path(args.output_dir)
    if not output_path.exists():
        output_path.mkdir(parents=True, exist_ok=True)

    i = 0
    j = 0
    for fn in tqdm(files):
        j += 1
        try:
            shape = read_step_file(str(fn))
            # render(shape, output_path.joinpath(f'{j:06d}' + ".png"), args.width, args.height)
            render(shape, output_path.joinpath(fn.stem[:6] + ".png"), args.width, args.height)
        except Exception as e:
            i += 1
            # raise e
            print(e)
            continue
    print("error number: ", i)
    print("total number: ", j)

if __name__ == "__main__":
    main()