import os
import trimesh
import gradio as gr
import matplotlib.pyplot as plt
import pyvista as pv
from pyvista import Plotter

# Install and start Xvfb
os.system("apt-get update && apt-get install -y xvfb libgl1-mesa-glx")
pv.start_xvfb()

def generate_cad_and_apdl(length, width, height):
    """
    Generate a 3D CAD model (box only) using Trimesh and output APDL script.
    """
    cad_model = trimesh.creation.box(extents=(length, width, height))
    cad_path = "/tmp/demo_model.stl"
    cad_model.export(cad_path)

    apdl_script = f"""
/prep7
block, 0, {length}, 0, {width}, 0, {height}
/mesh
esize, 10
vmesh, all
allsel, all
/solu
solve
/exit
"""
    apdl_path = "/tmp/demo_model_apdl.txt"
    with open(apdl_path, "w") as f:
        f.write(apdl_script)

    return cad_model, cad_path, apdl_path

def visualize_2d(cad_model):
    vertices = cad_model.vertices
    edges = cad_model.edges
    fig, ax = plt.subplots(figsize=(6, 6))
    for edge in edges:
        start, end = edge
        ax.plot(
            [vertices[start][0], vertices[end][0]],
            [vertices[start][1], vertices[end][1]],
            color="black",
            linewidth=2
        )
    ax.set_title("2D Projection (Top View)")
    ax.set_xlabel("X-axis")
    ax.set_ylabel("Y-axis")
    ax.grid(True)
    ax.set_aspect('equal')
    image_path = "/tmp/cad_2d_visualization.png"
    plt.savefig(image_path)
    plt.close(fig)
    return image_path

def visualize_3d(cad_model):
    vertices = cad_model.vertices
    faces = cad_model.faces
    faces_pv = []
    for face in faces:
        faces_pv.append(len(face))
        faces_pv.extend(face)
    mesh = pv.PolyData(vertices, faces_pv)
    plotter = Plotter(off_screen=True)
    plotter.add_mesh(mesh, color="lightblue", show_edges=True)
    plotter.set_background("white")
    plotter.view_isometric()
    image_path = "/tmp/cad_3d_visualization.png"
    plotter.screenshot(image_path)
    plotter.close()
    return image_path

def create_files_and_visualizations(length, width, height):
    cad_model, cad_file, apdl_file = generate_cad_and_apdl(length, width, height)
    cad_2d_path = visualize_2d(cad_model)
    cad_3d_path = visualize_3d(cad_model)
    return cad_file, apdl_file, cad_2d_path, cad_3d_path

with gr.Blocks() as app:
    with gr.Row():
        length = gr.Number(label="Length (mm)", value=100, precision=2)
        width = gr.Number(label="Width (mm)", value=50, precision=2)
        height = gr.Number(label="Height (mm)", value=20, precision=2)
    submit_button = gr.Button("Submit")
    with gr.Row():
        cad_file = gr.File(label="Download CAD Model (STL)")
        apdl_file = gr.File(label="Download APDL Script")
        cad_2d_image = gr.Image(label="2D Visualization (Top View)")
        cad_3d_image = gr.Image(label="3D Visualization")
    submit_button.click(
        fn=create_files_and_visualizations,
        inputs=[length, width, height],
        outputs=[cad_file, apdl_file, cad_2d_image, cad_3d_image]
    )

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