import streamlit as st
import tempfile
import subprocess
import os
import requests
import base64
import sys
import shutil
import re
from streamlit_ace import st_ace

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

# Preprocess script to ensure it has necessary imports and path handling
def preprocess_script(script_text, tmp_dir):
    # Check if 'import os' exists in the script
    if 'import os' not in script_text:
        script_text = 'import os\n' + script_text
    
    # Check if 'import math' exists in the script (commonly needed)
    if 'import math' not in script_text and 'math.' in script_text:
        script_text = 'import math\n' + script_text
    
    # Check if script already has file saving code
    if 'bpy.ops.wm.save_as_mainfile' not in script_text:
        # Add file saving code at the end
        script_text += f'''

# Save the .blend file
output_dir = os.environ.get('BLENDER_OUTPUT_DIR', os.getcwd())
blend_file_path = os.path.join(output_dir, "scene.blend")
bpy.ops.wm.save_as_mainfile(filepath=blend_file_path)
print(f"Saved blend file to: {{blend_file_path}}")
'''
    else:
        # Make sure existing save code uses absolute paths correctly
        script_text = re.sub(
            r'bpy\.ops\.wm\.save_as_mainfile\s*\(\s*filepath\s*=\s*["\']([^"\']+)["\']',
            'bpy.ops.wm.save_as_mainfile(filepath=os.path.join(os.environ.get("BLENDER_OUTPUT_DIR", os.getcwd()), "scene.blend")',
            script_text
        )
    
    return script_text

# 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 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
"""

# Sidebar: Blender script examples
script_examples = {
    "Earth with Texture": default_script,
    "Solar System": """import bpy
import math

# Clear existing objects
bpy.ops.object.select_all(action='SELECT')
bpy.ops.object.delete(use_global=False)

# Function to create a sphere (planet or sun)
def create_celestial_body(name, radius, location, color):
    # Create mesh sphere
    bpy.ops.mesh.primitive_uv_sphere_add(radius=radius, location=location)
    obj = bpy.context.active_object
    obj.name = name
    
    # Create material
    mat = bpy.data.materials.new(name + "_Material")
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get('Principled BSDF')
    bsdf.inputs['Base Color'].default_value = (*color, 1)
    bsdf.inputs['Roughness'].default_value = 0.5
    
    # Assign material
    if obj.data.materials:
        obj.data.materials[0] = mat
    else:
        obj.data.materials.append(mat)
    return obj

# Create the Sun (yellow)
sun = create_celestial_body('Sun', radius=2, location=(0,0,0), color=(1.0, 1.0, 0.0))

# Create an empty at origin to parent planets for orbiting
def create_orbit_empty(name):
    bpy.ops.object.empty_add(type='PLAIN_AXES', location=(0,0,0))
    empty = bpy.context.active_object
    empty.name = name
    return empty

# Planet definitions: name, radius, distance from sun, color, orbital period in frames
planets = [
    ('Mercury', 0.3, 4, (0.8, 0.5, 0.2), 88),
    ('Venus',   0.5, 6, (1.0, 0.8, 0.0), 224),
    ('Earth',   0.5, 8, (0.2, 0.4, 1.0), 365),
    ('Mars',    0.4, 10,(1.0, 0.3, 0.2), 687)
]

for name, radius, distance, color, period in planets:
    # Create orbit controller
    orbit_empty = create_orbit_empty(name + '_Orbit')
    
    # Create planet as child of orbit empty
    planet = create_celestial_body(name, radius, location=(distance, 0, 0), color=color)
    planet.parent = orbit_empty
    
    # Animate the orbit: rotate empty around Z
    orbit_empty.rotation_euler = (0, 0, 0)
    orbit_empty.keyframe_insert(data_path='rotation_euler', frame=1, index=2)
    orbit_empty.rotation_euler = (0, 0, math.radians(360))
    orbit_empty.keyframe_insert(data_path='rotation_euler', frame=period, index=2)
    
    # Set linear interpolation for smooth constant motion
    action = orbit_empty.animation_data.action
    fcurve = action.fcurves.find('rotation_euler', index=2)
    if fcurve:
        for kp in fcurve.keyframe_points:
            kp.interpolation = 'LINEAR'

# Set scene frames
bpy.context.scene.frame_start = 1
bpy.context.scene.frame_end = 687  # Use Mars' period as the end frame

# Add simple starfield as world background
world = bpy.context.scene.world
world.use_nodes = True
bg = world.node_tree.nodes['Background']
bg.inputs['Color'].default_value = (0, 0, 0, 1)

# Set up a better camera view
bpy.ops.object.camera_add(location=(0, -20, 15))
camera = bpy.context.active_object
camera.rotation_euler = (math.radians(55), 0, 0)
bpy.context.scene.camera = camera

# Add a sun light
bpy.ops.object.light_add(type='SUN', location=(10, -10, 10))
sun_light = bpy.context.active_object
sun_light.data.energy = 5.0

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

print('Solar system animation setup complete.')"""
}

# Add example selector
selected_example = st.sidebar.selectbox(
    "Load example script:",
    options=list(script_examples.keys())
)

# Sidebar: Blender script editor
st.sidebar.header("Blender Python Script")
script_text = st_ace(
    value=script_examples[selected_example],
    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) Preprocess and save user script
            processed_script = preprocess_script(script_text, tmp_dir)
            script_path = os.path.join(tmp_dir, "user_script.py")
            with open(script_path, "w") as f:
                f.write(processed_script)

            # 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 necessary variables
            env = os.environ.copy()
            env["TEXTURE_PATHS"] = ",".join(texture_paths)
            env["BLENDER_OUTPUT_DIR"] = tmp_dir  # Pass the tmp_dir to the Blender script
            
            # 5) Run Blender to build .blend file
            blend_path = os.path.join(tmp_dir, "scene.blend")
            
            with st.status("Running Blender to create scene...") as status:
                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
                    )
                    status.update(label="Blender scene created successfully", state="complete")
                    st.text_area("Blender output", r1.stdout + r1.stderr, height=150)
                except subprocess.TimeoutExpired:
                    st.error("Blender process timed out after 3 minutes.")
                    st.stop()
                except subprocess.CalledProcessError as e:
                    st.error(f"Blender build failed with error code {e.returncode}")
                    st.text_area("Error details", e.stdout + e.stderr, height=150)
                    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.info("This might be due to an issue with the script or Blender configuration.")
                st.stop()
                
            # 6) Export GLB with animation
            glb_path = os.path.join(tmp_dir, "animation.glb")
            
            with st.status("Exporting to GLB format...") as status:
                # Create a separate Python script for GLB export to avoid command line issues
                export_script = os.path.join(tmp_dir, "export_glb.py")
                with open(export_script, "w") as f:
                    f.write(f"""
import bpy
import os

# Load the blend file
bpy.ops.wm.open_mainfile(filepath=r'{blend_path}')

# Export to GLB
bpy.ops.export_scene.gltf(
    filepath=r'{glb_path}',
    export_format='GLB',
    export_animations=True
)
""")
                
                cmd2 = [blender_path, "--background", "--python", export_script]
                try:
                    r2 = subprocess.run(
                        cmd2,
                        cwd=tmp_dir,
                        env=env,
                        check=True,
                        capture_output=True,
                        text=True,
                        timeout=120  # 2 minute timeout
                    )
                    status.update(label="GLB export completed successfully", state="complete")
                    st.text_area("Export output", r2.stdout + r2.stderr, height=100)
                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()
                
                file_size = len(data) / (1024 * 1024)  # Size in MB
                st.success(f"Successfully created GLB file ({file_size:.1f} MB)")
                
                # Check if file isn't too large for embedding
                if file_size > 50:
                    st.warning("The GLB file is quite large. The viewer might be slow to load.")
                
                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 
        ar
        shadow-intensity="1"
        style="width:100%; height:600px; background-color: #f0f0f0;">
    </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.")
                if os.path.exists(blend_path):
                    st.info("The Blender file was created, but the GLB export failed.")
                
        finally:
            # Clean up - remove temporary directory
            try:
                shutil.rmtree(tmp_dir, ignore_errors=True)
            except Exception as e:
                st.warning(f"Failed to clean up temporary files: {str(e)}")

# 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 app automatically adds necessary imports (`import os`) and file saving code to your scripts,
    so you can focus on creating 3D content without worrying about environment details.
    """)
    
with st.expander("Script Tips"):
    st.markdown("""
    ### Tips for writing Blender scripts
    
    1. **Accessing textures**: Use the environment variable `TEXTURE_PATHS` to get the paths to uploaded textures
    
    ```python
    # Example code to load the first uploaded texture
    texture_paths = os.environ.get('TEXTURE_PATHS', '').split(',')
    if texture_paths and texture_paths[0]:
        image = bpy.data.images.load(texture_paths[0])
    ```
    
    2. **You don't need to add file saving code** - the app automatically adds it
    
    3. **Animations**: Make sure to set keyframes if you want your model to animate:
    
    ```python
    # Example animation (rotate object 360 degrees over 250 frames)
    obj.rotation_euler = (0, 0, 0)
    obj.keyframe_insert(data_path="rotation_euler", frame=1)
    obj.rotation_euler = (0, 0, math.radians(360))
    obj.keyframe_insert(data_path="rotation_euler", frame=250)
    ```
    """)