import trimesh
import gradio as gr
import matplotlib.pyplot as plt
import io
import base64
from PIL import Image

def generate_cad_model(base_length, base_width, base_height):
    """
    Generate a CAD model as a 3D mesh.
    """
    # Create a 3D box (CAD model representation)
    mesh = trimesh.creation.box(extents=(base_length, base_width, base_height))
    return mesh

def render_3d_cad_model(base_length, base_width, base_height):
    """
    Render a 3D CAD model for interactive viewing.
    """
    # Generate the 3D CAD model
    mesh = generate_cad_model(base_length, base_width, base_height)
    
    # Convert the mesh to GLB format for 3D visualization
    glb_data = io.BytesIO()
    mesh.export(glb_data, file_type="glb")
    glb_data.seek(0)
    return glb_data

def render_2d_cad_projection(base_length, base_width, base_height):
    """
    Generate a 2D projection of the CAD model and return it as an image.
    """
    # Generate the 3D CAD model
    mesh = generate_cad_model(base_length, base_width, base_height)
    
    # Create a 2D projection of the model
    projection = mesh.section(plane_origin=(0, 0, 0), plane_normal=(0, 0, 1))
    if projection is not None:
        projection_img = trimesh.util.grid_2d(projection.to_image(), tile=(1, 1))
        
        # Save the projection as an image
        buffer = io.BytesIO()
        plt.imshow(projection_img, cmap="gray")
        plt.axis("off")
        plt.savefig(buffer, format="png", bbox_inches="tight")
        buffer.seek(0)
        
        # Convert to Base64 for Gradio display
        encoded_image = base64.b64encode(buffer.getvalue()).decode("utf-8")
        return f"data:image/png;base64,{encoded_image}"
    else:
        return None

# Gradio Interface
def display_cad(base_length, base_width, base_height):
    """
    Render both 3D and 2D CAD outputs.
    """
    # Render 3D model
    glb_data = render_3d_cad_model(base_length, base_width, base_height)
    
    # Render 2D projection
    cad_image = render_2d_cad_projection(base_length, base_width, base_height)
    
    return glb_data, cad_image

# Gradio interface for 3D and 2D CAD visualization
demo = gr.Interface(
    fn=display_cad,
    inputs=[
        gr.Number(label="Base Length (mm)", value=10),
        gr.Number(label="Base Width (mm)", value=10),
        gr.Number(label="Base Height (mm)", value=5),
    ],
    outputs=[
        gr.Model3D(label="3D CAD Model"),  # Interactive 3D viewer
        gr.Image(label="2D CAD Projection"),  # Static 2D projection
    ],
    live=True,
)

if __name__ == "__main__":
    demo.launch()