import trimesh
import numpy as np

def convert_vrml_to_stl(input_file, output_file):
    # Carrega o arquivo VRML (WRL)
    mesh = trimesh.load(input_file, file_type='vrml')

    # Verifica se a malha possui informações de cor
    if not hasattr(mesh.visual, 'vertex_colors') or mesh.visual.vertex_colors is None:
        raise ValueError("A malha não possui informações de cores")

    # Extrai as informações de cor
    colors = mesh.visual.vertex_colors

    # Cria um novo arquivo STL com as informações de cor embutidas como comentários
    with open(output_file, 'w') as stl_file:
        stl_file.write("solid converted_vrml\n")
        
        for face in mesh.faces:
            vertices = mesh.vertices[face]
            face_color = colors[face].mean(axis=0)[:3]  # Média das cores dos vértices do triângulo
            color_comment = f"color {face_color[0]} {face_color[1]} {face_color[2]}"
            stl_file.write(f"facet normal 0 0 0\n")
            stl_file.write(f"  outer loop\n")
            for vertex in vertices:
                stl_file.write(f"    vertex {vertex[0]} {vertex[1]} {vertex[2]}\n")
            stl_file.write(f"  endloop\n")
            stl_file.write(f"endfacet\n")
            stl_file.write(f"  {color_comment}\n")
        
        stl_file.write("endsolid converted_vrml\n")

    print(f"Conversão concluída. Arquivo STL salvo em: {output_file}")

if __name__ == "__main__":
    input_file = 'montagem_completa.wrl'  # Substitua pelo caminho do seu arquivo VRML de entrada
    output_file = 'montagem_completa.stl'  # Substitua pelo caminho do seu arquivo STL de saída

    convert_vrml_to_stl(input_file, output_file)