import streamlit as st

# Attempt to import required libraries and handle missing dependencies
try:
    from rembg import remove
except ModuleNotFoundError:
    st.error("The `rembg` module requires `onnxruntime`. Please install it using `pip install onnxruntime` or `pip install onnxruntime-gpu` if you have a GPU.")

import io
import trimesh
from trimesh import transformations
from PIL import Image
from io import BytesIO
import numpy as np

# Set the title of the app
st.title('Image to 3D Asset Generator')

# Upload image file
uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])

# Shape options
shape_options = [
    "Box", "Sphere", "Cone", "Cylinder", "Torus", "Icosahedron", "Octahedron", "Dodecahedron",
    "Grid", "Plane"
]

# Select the shape
shape = st.selectbox("Choose a 3D Shape", shape_options)

# Transformation Options
scale_x = st.slider("Scale X", min_value=0.1, max_value=5.0, value=1.0)
scale_y = st.slider("Scale Y", min_value=0.1, max_value=5.0, value=1.0)
scale_z = st.slider("Scale Z", min_value=0.1, max_value=5.0, value=1.0)

rotation_x = st.slider("Rotate X (degrees)", min_value=0, max_value=360, value=0)
rotation_y = st.slider("Rotate Y (degrees)", min_value=0, max_value=360, value=0)
rotation_z = st.slider("Rotate Z (degrees)", min_value=0, max_value=360, value=0)

position_x = st.slider("Position X", min_value=-5.0, max_value=5.0, value=0.0)
position_y = st.slider("Position Y", min_value=-5.0, max_value=5.0, value=0.0)
position_z = st.slider("Position Z", min_value=-5.0, max_value=5.0, value=0.0)

# Color and Material options
color = st.color_picker("Choose a Color", "#00ff00")
file_format = st.selectbox("Choose Export Format", ["GLB", "OBJ", "STL"])

# Process the image and generate 3D asset
if uploaded_file is not None:
    try:
        # Display the image
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
        
        # Process the image for background removal (if `rembg` is available)
        if "remove" in globals():
            input_data = uploaded_file.read()
            output_data = remove(input_data)
            # Display the processed image (background removed)
            output_image = Image.open(BytesIO(output_data))
            st.image(output_image, caption="Processed Image (Background Removed)", use_column_width=True)
        else:
            st.warning("Image processing with background removal is unavailable. Please install `onnxruntime` and restart the app.")
        
        # Create the selected shape and apply transformations
        mesh = None
        if shape == "Box":
            mesh = trimesh.creation.box(extents=[1, 1, 1])
        elif shape == "Sphere":
            mesh = trimesh.creation.uv_sphere(radius=1.0)
        elif shape == "Cone":
            mesh = trimesh.creation.cone(radius=1.0, height=1.0)
        elif shape == "Cylinder":
            mesh = trimesh.creation.cylinder(radius=1.0, height=1.0)
        elif shape == "Torus":
            mesh = trimesh.creation.torus(radius=1.0, tube_radius=0.5)
        elif shape == "Icosahedron":
            mesh = trimesh.creation.icosahedron()
        elif shape == "Octahedron":
            mesh = trimesh.creation.octahedron()
        elif shape == "Dodecahedron":
            mesh = trimesh.creation.dodecahedron()
        elif shape == "Grid":
            mesh = trimesh.creation.box(extents=[5, 5, 0.1])
        elif shape == "Plane":
            mesh = trimesh.creation.box(extents=[5, 5, 0.01])
        else:
            st.error("Unsupported shape selected.")

        if mesh:
            # Apply transformations
            mesh.apply_scale([scale_x, scale_y, scale_z])
            rotation_matrix = transformations.euler_matrix(
                np.radians(rotation_x), np.radians(rotation_y), np.radians(rotation_z)
            )
            mesh.apply_transform(rotation_matrix)
            mesh.apply_translation([position_x, position_y, position_z])

            # Apply color
            material_color = np.array(
                [int(color[1:3], 16) / 255, int(color[3:5], 16) / 255, int(color[5:7], 16) / 255, 1.0]
            )
            mesh.visual.face_colors = (material_color * 255).astype(np.uint8)
            
            # Export the mesh
            with io.BytesIO() as buffer:
                mesh.export(buffer, file_type=file_format.lower())
                buffer.seek(0)
                st.download_button(
                    label=f"Download 3D Asset ({file_format})",
                    data=buffer,
                    file_name=f"output_model.{file_format.lower()}",
                    mime=f"model/{file_format.lower()}"
                )
    except Exception as e:
        st.error(f"An error occurred while generating the 3D asset: {e}")