import cadquery as cq
import os
import pyvista as pv

def generate_die(length, width, thickness):
    try:
        output_dir = "outputs"
        os.makedirs(output_dir, exist_ok=True)

        plate = cq.Workplane("XY").box(length, width, thickness)
        punch = cq.Workplane("XY").rect(10, 10).extrude(5).translate((length / 4, width / 4, thickness / 2))
        die = plate.cut(punch)

        filename = os.path.join(output_dir, "progressive_die.step")
        cq.exporters.export(die, filename)

        return filename
    except Exception as e:
        return f"Error generating die: {str(e)}"

def visualize_die(length, width, thickness):
    try:
        output_dir = "outputs"
        os.makedirs(output_dir, exist_ok=True)

        plate = cq.Workplane("XY").box(length, width, thickness)
        punch = cq.Workplane("XY").rect(10, 10).extrude(5).translate((length / 4, width / 4, thickness / 2))
        die = plate.cut(punch)

        stl_file = os.path.join(output_dir, "progressive_die.stl")
        cq.exporters.exportShape(die.val(), "STL", stl_file)

        pv.global_theme.off_screen = True
        mesh = pv.read(stl_file)
        plotter = pv.Plotter(off_screen=True)
        plotter.add_mesh(mesh, color="blue")
        screenshot = os.path.join(output_dir, "progressive_die_visualization.png")
        plotter.screenshot(screenshot)

        return screenshot
    except Exception as e:
        return f"Error visualizing die: {str(e)}"