import gradio as gr
import cadquery as cq
import os
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# G-code Generator Function
def process_cad_and_generate_gcode(cad_file, feed_rate, spindle_speed):
    """
    Processes the uploaded CAD file, generates basic G-code, and visualizes the 3D shape.
    """
    if not cad_file:
        return "No file provided. Please upload a valid CAD file.", None

    # Validate CAD file extension
    file_extension = os.path.splitext(cad_file.name)[1].lower()
    if file_extension not in [".step", ".stp", ".iges", ".igs"]:
        return "Unsupported file type. Please upload a STEP or IGES file.", None

    # Load the CAD file into CadQuery
    try:
        shape = cq.importers.importStep(cad_file.name)
    except Exception as e:
        return f"Error processing CAD file: {str(e)}", None

    # Generate G-code simulation (simplified linear G-code for demo purposes)
    gcode = []
    gcode.append("G21 ; Set units to millimeters")
    gcode.append("G90 ; Absolute positioning")
    gcode.append("G28 ; Move to home position")
    gcode.append(f"M3 S{spindle_speed} ; Start spindle")
    gcode.append(f"G01 Z-3 F{feed_rate} ; Move to depth of 3mm")
    gcode.append("G01 X0 Y0 F100 ; Start machining")
    gcode.append("G00 Z5 ; Retract tool")
    gcode.append("M5 ; Stop spindle")
    gcode.append("M30 ; End of program")

    # 3D Visualization using Matplotlib
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")

    # Extract bounding box for visualization
    bounding_box = shape.val().BoundingBox()
    X = [bounding_box.xmin, bounding_box.xmax, bounding_box.xmax, bounding_box.xmin, bounding_box.xmin]
    Y = [bounding_box.ymin, bounding_box.ymin, bounding_box.ymax, bounding_box.ymax, bounding_box.ymin]
    Z = [bounding_box.zmin] * 5

    verts = [list(zip(X, Y, Z))]
    ax.add_collection3d(Poly3DCollection(verts, color="lightblue", linewidths=1, edgecolor="r"))

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.set_title("3D Visualization of CAD File")

    # Save the figure
    output_3d_image = "3d_visualization.png"
    plt.savefig(output_3d_image)
    plt.close()

    return "\n".join(gcode), output_3d_image

# Gradio Interface
iface = gr.Interface(
    fn=process_cad_and_generate_gcode,
    inputs=[
        gr.File(label="Upload CAD File (STEP, IGES)"),
        gr.Slider(50, 500, step=10, label="Feed Rate (mm/min)", value=100),
        gr.Slider(500, 5000, step=500, label="Spindle Speed (RPM)", value=1000),
    ],
    outputs=[
        gr.Textbox(label="Generated G-code"),
        gr.Image(label="3D Visualization of Processed CAD File"),
    ],
    title="CAD to G-code Generator with 3D Visualization",
    description="Upload a CAD file (STEP, IGES), and this tool generates CNC G-code along with a 3D visualization of the part."
)

# Launch the Interface
if __name__ == "__main__":
    iface.launch(share=True)