from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.TCollection import TCollection_ExtendedString, TCollection_AsciiString
from OCC.Core.XCAFDoc import (
    XCAFDoc_DocumentTool_ShapeTool,
    XCAFDoc_DocumentTool_LayerTool,
    XCAFDoc_DocumentTool_VisMaterialTool,
    XCAFDoc_DocumentTool_ColorTool,
    XCAFDoc_VisMaterial,
    XCAFDoc_VisMaterialPBR
)
from OCC.Core.TopAbs import TopAbs_SOLID, TopAbs_FACE, TopAbs_SHAPE, TopAbs_COMPOUND, TopAbs_COMPSOLID, TopAbs_SHELL, \
    TopAbs_VERTEX, TopAbs_EDGE, TopAbs_WIRE
from OCC.Core.TColStd import TColStd_IndexedDataMapOfStringString
from OCC.Core.Message import Message_ProgressRange
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Extend.DataExchange import read_step_file_with_names_colors
# GLTF export
from OCC.Core.RWGltf import RWGltf_CafWriter, RWGltf_WriterTrsfFormat

from OCC.Core.Quantity import (
    Quantity_Color,
    Quantity_ColorRGBA,
    Quantity_TOC_RGB,
    Quantity_NOC_BLUE1
)
from OCC.Core.TCollection import TCollection_ExtendedString, TCollection_AsciiString
from OCC.Display.SimpleGui import init_display

display, start_display, add_menu, add_function_to_menu = init_display()

# create a document
doc = TDocStd_Document(TCollection_ExtendedString("pythonocc-doc"))
shape_tool = XCAFDoc_DocumentTool_ShapeTool(doc.Main())
layer_tool = XCAFDoc_DocumentTool_LayerTool(doc.Main())
color_tool=XCAFDoc_DocumentTool_ColorTool(doc.Main())
vm_tool=XCAFDoc_DocumentTool_VisMaterialTool(doc.Main())

filename = "as1_pe_203.stp"
shapes_labels_colors = read_step_file_with_names_colors(filename)

id=0
for a_shape in shapes_labels_colors:
    s, c = shapes_labels_colors[a_shape]

    # Test Set Material
    vismat= XCAFDoc_VisMaterial()
    visMatPbr=XCAFDoc_VisMaterialPBR()
    visMatPbr.BaseColor=Quantity_ColorRGBA(Quantity_Color(c.Red(),c.Green(),c.Blue(),Quantity_TOC_RGB))
    visMatPbr.IsDefined=True
    vismat.SetPbrMaterial(visMatPbr)
    matlabel=vm_tool.AddMaterial(vismat,TCollection_AsciiString("mymat"+str(id)))
    vm_tool.SetShapeMaterial(a_shape,matlabel)
    
    # Test Add Color 
    color_tool.AddColor(Quantity_Color(c.Red(),c.Green(),c.Blue(),Quantity_TOC_RGB))
    color_tool.SetInstanceColor(a_shape,0,Quantity_Color(c.Red(),c.Green(),c.Blue(),Quantity_TOC_RGB),True)
    color_tool.SetInstanceColor(a_shape,1,Quantity_Color(c.Red(),c.Green(),c.Blue(),Quantity_TOC_RGB),True)
    color_tool.SetInstanceColor(a_shape,2,Quantity_Color(c.Red(),c.Green(),c.Blue(),Quantity_TOC_RGB),True)
    
    msh_algo = BRepMesh_IncrementalMesh()
    msh_algo.SetShape(a_shape)
    msh_algo.Perform()
    shape_tool.AddShape(a_shape)
    display.DisplayShape(a_shape)

    id+=1
    print(id)

start_display()
# GLTF options
a_format = RWGltf_WriterTrsfFormat.RWGltf_WriterTrsfFormat_Compact
force_uv_export = True

# metadata
a_file_info = TColStd_IndexedDataMapOfStringString()
a_file_info.Add(
    TCollection_AsciiString("Authors"), TCollection_AsciiString("pythonocc")
)

#
# Binary export
#
binary = True
binary_rwgltf_writer = RWGltf_CafWriter(TCollection_AsciiString("box.glb"), binary)
binary_rwgltf_writer.SetTransformationFormat(a_format)
binary_rwgltf_writer.SetForcedUVExport(force_uv_export)
pr = Message_ProgressRange()  # this is required
binary_rwgltf_writer.Perform(doc, a_file_info, pr)

#
# Ascii export
#
binary = False
ascii_rwgltf_writer = RWGltf_CafWriter(TCollection_AsciiString("box.gltf"), binary)
ascii_rwgltf_writer.SetTransformationFormat(a_format)
ascii_rwgltf_writer.SetForcedUVExport(force_uv_export)
pr = Message_ProgressRange()  # this is required
ascii_rwgltf_writer.Perform(doc, a_file_info, pr)


class RWGltf_CafWriterNew(RWGltf_CafWriter):
    def writeMaterials():
        pass