import cadquery as cq
import numpy as np
import pyvista as pv


def visualize_step(
    step_file  = "file.step",
    color      = "silver",
    background = "white",
    tolerance  = 0.1,
):
    # 1. Load the STEP file
    print("Loading STEP file...")
    shape = cq.importers.importStep(step_file).val()

    # 2. Convert the shape into a triangle mesh
    print("Tessellating...")
    raw_vertices, triangles = shape.tessellate(tolerance)

    vertices  = np.array([[v.x, v.y, v.z] for v in raw_vertices], dtype=np.float64)
    tri_array = np.array(triangles, dtype=np.int32)

    # PyVista needs faces in the format [3, i0, i1, i2,  3, i0, i1, i2, ...]
    n     = len(tri_array)
    faces = np.column_stack([np.full(n, 3, dtype=np.int32), tri_array]).ravel()

    mesh = pv.PolyData(vertices, faces)
    mesh.compute_normals(inplace=True)  # makes shading look nicer

    # 3. Show the model
    print("Opening viewer...")
    plotter = pv.Plotter(title="CAD Viewer")
    plotter.set_background(background)

    # Shaded solid surface
    plotter.add_mesh(mesh, color=color, smooth_shading=True, show_edges=False)

    # Black edges along sharp corners only
    edges = mesh.extract_feature_edges(feature_angle=20)
    if edges.n_points > 0:
        plotter.add_mesh(edges, color="black", line_width=2)

    plotter.add_axes()
    plotter.reset_camera()
    plotter.show()


# ── Call it here ───────────────────────────────────────────────────────────────
visualize_step("file.step")