import streamlit as st
import numpy as np
import pyvista as pv
from stl import mesh
import tempfile
import os
import trimesh

# Page configuration
st.set_page_config(page_title="AR CAD Viewer", layout="wide")

# Title
st.title("Augmented Reality CAD Viewer")
st.write("Upload your CAD models, preview them, and convert them to AR-compatible formats.")

# File uploader for CAD files
uploaded_file = st.file_uploader("Upload a CAD file (.stl format)", type=["stl"])

if uploaded_file is not None:
    # Save uploaded file to a temporary location
    with tempfile.NamedTemporaryFile(delete=False, suffix=".stl") as temp_file:
        temp_file.write(uploaded_file.read())
        temp_filepath = temp_file.name

    # Load the STL file using pyvista for interactive 3D rendering
    st.subheader("3D Model Preview")
    
    # Load the STL using pyvista
    mesh_data = pv.read(temp_filepath)
    
    # Create the plotter and add the mesh to it
    plotter = pv.Plotter(off_screen=True)
    plotter.add_mesh(mesh_data, color="cyan")
    plotter.view_isometric()
    
    # Show the plot in the Streamlit app
    st.pydeck_chart(plotter.show())

    # Convert to glTF for AR
    st.subheader("Convert to AR Format")
    st.write("Click below to convert the STL file into an AR-compatible `.glTF` format.")

    if st.button("Convert to glTF"):
        # Use Trimesh for conversion
        trimesh_mesh = trimesh.load(temp_filepath)
        gltf_filepath = temp_filepath.replace(".stl", ".gltf")
        trimesh.exchange.gltf.export_gltf(trimesh_mesh, gltf_filepath)

        # Display download link for glTF
        with open(gltf_filepath, "rb") as gltf_file:
            st.download_button(
                label="Download glTF Model",
                data=gltf_file,
                file_name="model.gltf",
                mime="model/gltf+json"
            )

        # Clean up temporary files
        os.remove(gltf_filepath)
        os.remove(temp_filepath)

    # AR Hosting Info
    st.subheader("Host in AR")
    st.markdown(
        """
        Upload your `.glTF` file to [GitHub Pages](https://pages.github.com/) or use a service like [Google Poly](https://poly.google.com/) (now integrated into other platforms).
        
        For AR.js:
        - Host the `.glTF` file on a web server.
        - Create an AR marker or use markerless AR for viewing.

        Example: `<a-scene><a-entity gltf-model="model.gltf"></a-entity></a-scene>`
        """
    )