# Updated app.py
import streamlit as st
import tempfile
import subprocess
import os
import requests
import base64
import sys
from streamlit_ace import st_ace
import shutil

st.set_page_config(page_title="Blender 3D Viewer", layout="wide")
st.title("🌍 Blender Script → 3D Viewer")

# Find Blender executable
def find_blender():
    # Common paths to check
    paths_to_check = [
        "blender",  # default PATH
        "/usr/bin/blender",
        "/opt/blender/blender"
    ]
    
    for path in paths_to_check:
        try:
            result = subprocess.run([path, "--version"], 
                                   capture_output=True, 
                                   text=True,
                                   timeout=5)
            if result.returncode == 0:
                st.success(f"Found Blender: {result.stdout.strip()}")
                return path
        except:
            continue
    
    return None

# Check if Blender is available at startup
blender_path = find_blender()
if not blender_path:
    st.error("❌ Blender not found! The app requires Blender to be installed.")
    st.info("Please check the Spaces configuration for proper apt.txt setup.")
    # Continue anyway to show the interface

# Hardcoded default texture URLs (with fallbacks)
DEFAULT_TEXTURE_URLS = [
    "https://eoimages.gsfc.nasa.gov/images/imagerecords/57000/57730/land_ocean_ice_2048.jpg",
    "https://svs.gsfc.nasa.gov/vis/a000000/a002900/a002915/bluemarble-2048.png"
]

# Sidebar: Blender script editor with example script
default_script = """
import bpy
import os
import math

# Clear the scene
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete()

# Create a UV Sphere for Earth
bpy.ops.mesh.primitive_uv_sphere_add(radius=1, location=(0, 0, 0))
earth = bpy.context.active_object
earth.name = 'Earth'

# Apply material to Earth
mat = bpy.data.materials.new(name="EarthMaterial")
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links

# Clear default nodes
for node in nodes:
    nodes.remove(node)

# Create texture node
texture_node = nodes.new(type='ShaderNodeTexImage')

# Get texture path from environment variable if available
texture_paths = os.environ.get('TEXTURE_PATHS', '').split(',')
if texture_paths and texture_paths[0]:
    texture_node.image = bpy.data.images.load(texture_paths[0])
    
# Add Principled BSDF node
principled = nodes.new(type='ShaderNodeBsdfPrincipled')
principled.inputs['Specular'].default_value = 0.1
principled.inputs['Roughness'].default_value = 0.8

# Add Output node
output = nodes.new(type='ShaderNodeOutputMaterial')

# Link nodes
links.new(texture_node.outputs['Color'], principled.inputs['Base Color'])
links.new(principled.outputs['BSDF'], output.inputs['Surface'])

# Assign material to Earth
if earth.data.materials:
    earth.data.materials[0] = mat
else:
    earth.data.materials.append(mat)

# Add a simple animation - rotation
earth.rotation_euler = (0, 0, 0)
earth.keyframe_insert(data_path="rotation_euler", frame=1)

# Rotate 360 degrees on Z axis
earth.rotation_euler = (0, 0, math.radians(360))
earth.keyframe_insert(data_path="rotation_euler", frame=250)

# Set animation interpolation to linear
for fc in earth.animation_data.action.fcurves:
    for kf in fc.keyframe_points:
        kf.interpolation = 'LINEAR'

# Setup camera
bpy.ops.object.camera_add(location=(0, -3, 0))
camera = bpy.context.active_object
camera.rotation_euler = (math.radians(90), 0, 0)
bpy.context.scene.camera = camera

# Setup lighting
bpy.ops.object.light_add(type='SUN', location=(5, -5, 5))
light = bpy.context.active_object
light.data.energy = 2.0

# Set render settings
bpy.context.scene.render.engine = 'CYCLES'
bpy.context.scene.cycles.device = 'CPU'
bpy.context.scene.render.film_transparent = True

# Save the .blend file in the current directory
blend_file_path = bpy.path.abspath("//scene.blend")
bpy.ops.wm.save_as_mainfile(filepath=blend_file_path)
"""

# Sidebar: Blender script editor
st.sidebar.header("Blender Python Script")
script_text = st_ace(
    value=default_script,
    placeholder="Paste your Blender Python script here...",
    language="python",
    theme="monokai",
    key="ace",
    min_lines=20,
    max_lines=100,
)

# Sidebar: texture uploads
st.sidebar.header("Texture Uploads (JPG/PNG)")
uploaded_textures = st.sidebar.file_uploader(
    "Upload one or more textures", type=["jpg", "jpeg", "png"], accept_multiple_files=True
)

# Sidebar: custom Python libraries
st.sidebar.header("Custom Python Libraries")
custom_packages = st.sidebar.text_area(
    "List pip packages (one per line)",
    height=100
)

# Main action
if st.sidebar.button("Run & Export GLB"):
    if not blender_path:
        st.error("Cannot proceed: Blender is not available")
        st.stop()
        
    if not script_text or not script_text.strip():
        st.error("Please provide a valid Blender Python script.")
        st.stop()
    
    with st.spinner("Processing your 3D scene..."):
        # Use a more robust temporary directory approach
        tmp_dir = tempfile.mkdtemp(prefix="blender_app_")
        try:
            # 1) Save user script
            script_path = os.path.join(tmp_dir, "user_script.py")
            with open(script_path, "w") as f:
                f.write(script_text)

            # 2) Collect textures
            texture_paths = []
            if uploaded_textures:
                for idx, upload in enumerate(uploaded_textures):
                    ext = os.path.splitext(upload.name)[1]
                    path = os.path.join(tmp_dir, f"texture_{idx}{ext}")
                    with open(path, "wb") as tf:
                        tf.write(upload.read())
                    texture_paths.append(path)
            else:
                # Try multiple default textures if some fail
                for url in DEFAULT_TEXTURE_URLS:
                    try:
                        r = requests.get(url, timeout=10)
                        r.raise_for_status()
                        ext = os.path.splitext(url)[-1] or ".jpg"
                        path = os.path.join(tmp_dir, f"default{ext}")
                        with open(path, "wb") as tf:
                            tf.write(r.content)
                        texture_paths.append(path)
                        # If we got one texture successfully, that's enough
                        break
                    except Exception as e:
                        st.warning(f"Could not download texture {url}: {str(e)}")

            # Display texture file information
            if texture_paths:
                st.subheader("Texture Files")
                for tp in texture_paths:
                    st.code(os.path.basename(tp))
            else:
                st.warning("No textures were loaded. Your 3D model may appear without textures.")

            # 3) Install custom Python libraries if specified
            if custom_packages and custom_packages.strip():
                pkgs = [l.strip() for l in custom_packages.splitlines() if l.strip()]
                if pkgs:
                    st.info(f"Installing: {', '.join(pkgs)}")
                    try:
                        # Use Python executable from current environment for pip
                        pip_cmd = [sys.executable, "-m", "pip", "install", "--user"] + pkgs
                        pip_res = subprocess.run(
                            pip_cmd,
                            check=True,
                            capture_output=True,
                            text=True,
                            timeout=120  # 2 minute timeout
                        )
                        st.text_area("pip install output", pip_res.stdout + pip_res.stderr, height=100)
                    except Exception as e:
                        st.warning(f"pip install failed: {str(e)}")

            # 4) Prepare environment with minimal variables
            env = os.environ.copy()
            env["TEXTURE_PATHS"] = ",".join(texture_paths)
            
            # 5) Run Blender to build .blend file
            blend_path = os.path.join(tmp_dir, "scene.blend")
            
            with st.status("Running Blender to create scene..."):
                cmd1 = [blender_path, "--background", "--python", script_path]
                try:
                    r1 = subprocess.run(
                        cmd1,
                        cwd=tmp_dir,
                        env=env,
                        check=True,
                        capture_output=True,
                        text=True,
                        timeout=180  # 3 minute timeout
                    )
                    st.code(r1.stdout.split("\n")[-10:], language="text")  # Show just the last few lines
                except subprocess.TimeoutExpired:
                    st.error("Blender process timed out after 3 minutes.")
                    st.stop()
                except Exception as e:
                    st.error(f"Blender build failed: {str(e)}")
                    st.stop()
            
            # Check if blend file was created
            if not os.path.exists(blend_path):
                st.error("Blender did not create the expected scene.blend file.")
                st.stop()
                
            # 6) Export GLB with animation
            glb_path = os.path.join(tmp_dir, "animation.glb")
            
            with st.status("Exporting to GLB format..."):
                expr = (
                    "import bpy;"
                    "bpy.ops.wm.open_mainfile(filepath=r'" + blend_path + "');"
                    "bpy.ops.export_scene.gltf(filepath=r'" + glb_path + 
                    "', export_format='GLB', export_animations=True)"
                )
                
                cmd2 = [blender_path, "--background", "--python-expr", expr]
                try:
                    r2 = subprocess.run(
                        cmd2,
                        cwd=tmp_dir,
                        env=env,
                        check=True,
                        capture_output=True,
                        text=True,
                        timeout=120  # 2 minute timeout
                    )
                except Exception as e:
                    st.error(f"GLB export failed: {str(e)}")
                    st.stop()

            # 7) Embed GLB inline if it exists
            if os.path.exists(glb_path):
                with open(glb_path, 'rb') as f:
                    data = f.read()
                
                b64 = base64.b64encode(data).decode()
                
                # Display the 3D model viewer
                st.subheader("3D Model Viewer")
                html = f"""
    <script type="module" src="https://unpkg.com/@google/model-viewer@latest/dist/model-viewer.min.js"></script>
    <model-viewer 
        src="data:model/gltf-binary;base64,{b64}" 
        alt="3D Model" 
        auto-rotate 
        camera-controls 
        style="width:100%; height:600px;">
    </model-viewer>
    """
                st.components.v1.html(html, height=650)
                
                # Add download button
                st.download_button(
                    "⬇️ Download GLB File", 
                    data, 
                    file_name="animation.glb",
                    mime="model/gltf-binary"
                )
            else:
                st.error("GLB file was not generated successfully.")
                
        finally:
            # Clean up - remove temporary directory
            try:
                shutil.rmtree(tmp_dir, ignore_errors=True)
            except:
                pass

# Add helpful information at the bottom
st.divider()
with st.expander("About this app"):
    st.markdown("""
    **Blender 3D Viewer App** lets you:
    
    - Write Blender Python scripts directly in your browser
    - Upload custom textures for your 3D models
    - Generate and visualize 3D models with animations
    - Download the results as GLB files for use in other applications
    
    The example script creates a rotating Earth with the uploaded texture.
    """)
    
    st.info("If the app isn't working, make sure 'blender' is properly installed in the Spaces environment.")