Update app.py

app.py

@@ -6,7 +6,8 @@ import json
 import pandas as pd
 import uuid
 from PIL import Image, ImageDraw # For minimap (using buttons instead now)
-from streamlit_js_eval import streamlit_js_eval, sync # For JS communication
+# Corrected import: removed 'sync'
+from streamlit_js_eval import streamlit_js_eval # For JS communication
 
 # --- Constants ---
 SAVE_DIR = "saved_worlds"
@@ -22,23 +23,25 @@ os.makedirs(SAVE_DIR, exist_ok=True)
 def load_plot_metadata():
     """Scans save dir, sorts plots, calculates metadata."""
     plots = []
+    # Ensure consistent sorting, e.g., alphabetically which often aligns with plot_001, plot_002 etc.
     plot_files = sorted([f for f in os.listdir(SAVE_DIR) if f.endswith(".csv")])
 
     current_x_offset = 0.0
     for i, filename in enumerate(plot_files):
         # Extract name - assumes format like 'plot_001_MyName.csv' or just 'plot_001.csv'
         parts = filename[:-4].split('_') # Remove .csv and split by underscore
-        plot_id = parts[0] # Assume first part is ID/order
-        plot_name = " ".join(parts[1:]) if len(parts) > 1 else f"Plot {i+1}"
+        plot_id = filename[:-4] # Use filename (without ext) as a unique ID for now
+        plot_name = " ".join(parts[1:]) if len(parts) > 1 else f"Plot {i+1}" # Try to extract name
 
         plots.append({
-            'id': plot_id,
+            'id': plot_id, # Use filename as ID
             'name': plot_name,
             'filename': filename,
             'x_offset': current_x_offset
         })
         current_x_offset += PLOT_WIDTH
-    return plots, current_x_offset # Return plots and the next available offset
+    # Also return the offset where the next plot would start
+    return plots, current_x_offset
 
 def load_plot_objects(filename, x_offset):
     """Loads objects from a CSV, applying the plot's x_offset."""
@@ -46,22 +49,31 @@ def load_plot_objects(filename, x_offset):
     objects = []
     try:
         df = pd.read_csv(file_path)
-        if not all(col in df.columns for col in CSV_COLUMNS):
-             st.warning(f"CSV '{filename}' missing expected columns. Skipping.")
+        # Check if required columns exist, handle gracefully if not
+        if not all(col in df.columns for col in ['type', 'pos_x', 'pos_y', 'pos_z']):
+             st.warning(f"CSV '{filename}' missing essential columns (type, pos_x/y/z). Skipping.")
              return []
+        # Ensure optional columns default to something sensible if missing
+        if 'obj_id' not in df.columns: df['obj_id'] = [str(uuid.uuid4()) for _ in range(len(df))]
+        if 'rot_x' not in df.columns: df['rot_x'] = 0.0
+        if 'rot_y' not in df.columns: df['rot_y'] = 0.0
+        if 'rot_z' not in df.columns: df['rot_z'] = 0.0
+        if 'rot_order' not in df.columns: df['rot_order'] = 'XYZ'
+
 
         for _, row in df.iterrows():
             obj_data = row.to_dict()
-            # Apply world offset
+            # Apply world offset accumulated during loading
             obj_data['pos_x'] += x_offset
             objects.append(obj_data)
         return objects
     except FileNotFoundError:
+        # This shouldn't happen if called via load_plot_metadata results, but handle anyway
         st.error(f"File not found during object load: {filename}")
         return []
     except pd.errors.EmptyDataError:
-         st.info(f"Plot file '{filename}' is empty.")
-         return [] # Empty file is valid
+         # An empty file is valid, represents an empty plot
+         return []
     except Exception as e:
         st.error(f"Error loading objects from {filename}: {e}")
         return []
@@ -71,29 +83,46 @@ def save_plot_data(filename, objects_data_list, plot_x_offset):
     """Saves object data list to a new CSV file, making positions relative."""
     file_path = os.path.join(SAVE_DIR, filename)
     relative_objects = []
+    # Ensure objects_data_list is actually a list
+    if not isinstance(objects_data_list, list):
+        st.error("Invalid data format received for saving (expected a list).")
+        print("Invalid save data:", objects_data_list) # Log for debugging
+        return False
+
     for obj in objects_data_list:
-        # Ensure required fields exist before proceeding
-        if not all(k in obj for k in ['position', 'rotation', 'type', 'obj_id']):
-            print(f"Skipping malformed object during save: {obj}")
+        # Validate incoming object structure more carefully
+        pos = obj.get('position', {})
+        rot = obj.get('rotation', {})
+        obj_type = obj.get('type', 'Unknown')
+        obj_id = obj.get('obj_id', str(uuid.uuid4())) # Generate ID if missing from JS
+
+        if not all(k in pos for k in ['x', 'y', 'z']) or obj_type == 'Unknown':
+            print(f"Skipping malformed object during save prep: {obj}")
             continue
 
         relative_obj = {
-            'obj_id': obj.get('obj_id', str(uuid.uuid4())), # Generate ID if missing
-            'type': obj.get('type', 'Unknown'),
-            'pos_x': obj['position'].get('x', 0.0) - plot_x_offset, # Make relative
-            'pos_y': obj['position'].get('y', 0.0),
-            'pos_z': obj['position'].get('z', 0.0),
-            'rot_x': obj['rotation'].get('_x', 0.0),
-            'rot_y': obj['rotation'].get('_y', 0.0),
-            'rot_z': obj['rotation'].get('_z', 0.0),
-            'rot_order': obj['rotation'].get('_order', 'XYZ')
+            'obj_id': obj_id,
+            'type': obj_type,
+            'pos_x': pos.get('x', 0.0) - plot_x_offset, # Make relative to plot start
+            'pos_y': pos.get('y', 0.0),
+            'pos_z': pos.get('z', 0.0),
+            'rot_x': rot.get('_x', 0.0),
+            'rot_y': rot.get('_y', 0.0),
+            'rot_z': rot.get('_z', 0.0),
+            'rot_order': rot.get('_order', 'XYZ')
         }
         relative_objects.append(relative_obj)
 
     try:
-        df = pd.DataFrame(relative_objects, columns=CSV_COLUMNS)
-        df.to_csv(file_path, index=False)
-        st.success(f"Saved plot data to {filename}")
+        # Only save if there are objects to save
+        if relative_objects:
+            df = pd.DataFrame(relative_objects, columns=CSV_COLUMNS)
+            df.to_csv(file_path, index=False)
+            st.success(f"Saved {len(relative_objects)} objects to {filename}")
+        else:
+            # Create an empty file with headers if nothing new was placed
+            pd.DataFrame(columns=CSV_COLUMNS).to_csv(file_path, index=False)
+            st.info(f"Saved empty plot file: {filename}")
         return True
     except Exception as e:
         st.error(f"Failed to save plot data to {filename}: {e}")
@@ -110,15 +139,16 @@ if 'selected_object' not in st.session_state:
     st.session_state.selected_object = 'None'
 if 'new_plot_name' not in st.session_state:
     st.session_state.new_plot_name = ""
-if 'save_request_data' not in st.session_state: # To store data back from JS
-    st.session_state.save_request_data = None
+# Use a more descriptive key for clarity
+if 'js_save_data_result' not in st.session_state:
+    st.session_state.js_save_data_result = None
 
 # --- Load Plot Metadata ---
 # Cached function returns list of plots and the next starting x_offset
 plots_metadata, next_plot_x_offset = load_plot_metadata()
 
 # --- Load ALL Objects for Rendering ---
-# This could be slow with many plots! Consider optimization later.
+# This could be slow with many plots!
 all_initial_objects = []
 for plot in plots_metadata:
     all_initial_objects.extend(load_plot_objects(plot['filename'], plot['x_offset']))
@@ -131,18 +161,20 @@ with st.sidebar:
     st.caption("Click to teleport player to the start of a plot.")
 
     # Use columns for a horizontal button layout if desired
-    cols = st.columns(3)
+    max_cols = 3 # Adjust number of columns
+    cols = st.columns(max_cols)
     col_idx = 0
     for plot in plots_metadata:
         # Use an emoji + name for the button
-        button_label = f"➡️ {plot.get('name', plot['id'])}"
-        if cols[col_idx].button(button_label, key=f"nav_{plot['id']}"):
+        button_label = f"➡️ {plot.get('name', plot['id'])}" # Fallback to id if name missing
+        # Use plot filename (unique) as key
+        if cols[col_idx].button(button_label, key=f"nav_{plot['filename']}"):
             # Send command to JS to move the player
             target_x = plot['x_offset']
             streamlit_js_eval(js_code=f"teleportPlayer({target_x});")
-            # No rerun needed here, JS handles the move
+            # No rerun needed here, JS handles the move instantly
 
-        col_idx = (col_idx + 1) % 3 # Cycle through columns
+        col_idx = (col_idx + 1) % max_cols # Cycle through columns
 
     st.markdown("---")
 
@@ -151,126 +183,140 @@ with st.sidebar:
     object_types = ["None", "Simple House", "Tree", "Rock", "Fence Post"]
     current_object_index = 0
     try:
+        # Ensure robustness if selected_object is somehow not in list
         current_object_index = object_types.index(st.session_state.selected_object)
     except ValueError:
-        st.session_state.selected_object = "None" # Reset if invalid
+        st.session_state.selected_object = "None" # Reset to default
+        current_object_index = 0
 
     selected_object_type_widget = st.selectbox(
         "Select Object:",
         options=object_types,
         index=current_object_index,
-        key="selected_object_widget"
+        key="selected_object_widget" # Use a distinct key for the widget
     )
+    # Update session state only if the widget's value actually changes
     if selected_object_type_widget != st.session_state.selected_object:
          st.session_state.selected_object = selected_object_type_widget
-         # No rerun needed just for selection, JS will read state on next placement
+         # Trigger a targeted JS update if needed, or rely on injection next cycle
+         # No rerun needed just for selection if JS reads global state on demand
 
 
     st.markdown("---")
 
     # --- Saving ---
     st.header("Save New Plot")
+    # Ensure text input reflects current state value
     st.session_state.new_plot_name = st.text_input(
         "Name for New Plot:",
         value=st.session_state.new_plot_name,
-        placeholder="My Awesome Creation"
+        placeholder="My Awesome Creation",
+        key="new_plot_name_input" # Use distinct key
     )
 
-    if st.button("💾 Save Current Work as New Plot"):
-        # 1. Trigger JS function to get data
-        # This function `getSaveData` should be defined in index.html
-        # It should collect data for objects considered "new"
-        # and return them as a JSON string.
-        # The positions should be WORLD positions.
-        js_get_data_code = "getSaveData();"
-        # Use streamlit_js_eval to run JS and get data back into session_state
-        st.session_state.save_request_data = streamlit_js_eval(
-             js_code=js_get_data_code,
-             key="get_save_data" # Unique key for this eval call
-             # Removed want_output=True, value goes to key if specified
+    if st.button("💾 Save Current Work as New Plot", key="save_button"):
+        # 1. Trigger JS function `getSaveData()` defined in index.html
+        # This function collects data for newly placed objects and returns a JSON string.
+        # The key argument ('js_save_processor') stores the JS result in session_state.
+        streamlit_js_eval(
+             js_code="getSaveData();",
+             key="js_save_processor" # Store result under this key
         )
-        # Need to trigger a rerun to process the data AFTER it arrives
-        # Using sync() might help ensure the value is available immediately after? Let's try without first.
-        st.rerun() # Rerun to process the returned data below
+        # Small delay MAY sometimes help ensure the value is set before rerun, but usually not needed
+        # import time
+        # time.sleep(0.1)
+        st.rerun() # Rerun to process the result in the next step
 
 
-# --- Process Save Data (if received from JS via key) ---
-# This runs AFTER the rerun triggered by the save button
-save_data_from_js = st.session_state.get("save_request_data", None)
+# --- Process Save Data (if received from JS via the key) ---
+# Check the session state key set by the streamlit_js_eval call
+save_data_from_js = st.session_state.get("js_save_processor", None)
 
-if save_data_from_js:
+if save_data_from_js is not None: # Process only if data is present
     st.info("Received save data from client...")
+    save_processed_successfully = False
     try:
-        objects_to_save = json.loads(save_data_from_js) # Parse JSON string from JS
+        # Ensure data is treated as a string before loading json
+        if isinstance(save_data_from_js, str):
+            objects_to_save = json.loads(save_data_from_js)
+        else:
+             # Handle case where it might already be parsed by chance (less likely)
+             objects_to_save = save_data_from_js
 
-        if isinstance(objects_to_save, list) and len(objects_to_save) > 0:
+        # Proceed only if we have a list (even an empty one is ok now)
+        if isinstance(objects_to_save, list):
             # Determine filename for the new plot
-            new_plot_index = len(plots_metadata) + 1
-            plot_name_sanitized = "".join(c for c in st.session_state.new_plot_name if c.isalnum() or c in (' ', '_')).rstrip() or f"Plot_{new_plot_index}"
-            new_filename = f"plot_{new_plot_index:03d}_{plot_name_sanitized.replace(' ','_')}.csv"
+            new_plot_index = len(plots_metadata) # 0-based index -> number of plots
+            # Sanitize name: replace spaces, keep only alphanumeric/underscore
+            plot_name_sanitized = "".join(c for c in st.session_state.new_plot_name if c.isalnum() or c in (' ')).strip().replace(' ', '_')
+            if not plot_name_sanitized: # Ensure there is a name part
+                plot_name_sanitized = f"Plot_{new_plot_index + 1}"
+
+            new_filename = f"plot_{new_plot_index:03d}_{plot_name_sanitized}.csv"
 
             # Save the data, converting world coords to relative coords inside the func
             save_ok = save_plot_data(new_filename, objects_to_save, next_plot_x_offset)
 
             if save_ok:
-                # IMPORTANT: Clear the plot metadata cache so it reloads with the new file
+                # Clear the plot metadata cache so it reloads with the new file
                 load_plot_metadata.clear()
-                # Clear the JS data request state
-                st.session_state.save_request_data = None
-                # Reset the new plot name field
+                # Reset the new plot name field for next time
                 st.session_state.new_plot_name = ""
-                # Reset newly placed objects in JS? Needs another call.
+                # Reset newly placed objects in JS AFTER successful save
                 try:
+                     # This call tells JS to clear its internal 'newlyPlacedObjects' array
                      streamlit_js_eval(js_code="resetNewlyPlacedObjects();", key="reset_js_state")
                 except Exception as js_e:
                      st.warning(f"Could not reset JS state after save: {js_e}")
 
                 st.success(f"New plot '{plot_name_sanitized}' saved!")
-                # Rerun AGAIN to reload plots list, redraw minimap, and update JS injection
-                st.rerun()
+                save_processed_successfully = True
             else:
                 st.error("Failed to save plot data to file.")
-                # Clear the request data even if save failed to prevent retry loop
-                st.session_state.save_request_data = None
 
-        elif isinstance(objects_to_save, list) and len(objects_to_save) == 0:
-            st.warning("Nothing new to save.")
-            st.session_state.save_request_data = None # Clear request
         else:
-             st.error(f"Received invalid save data format from client: {type(objects_to_save)}")
-             st.session_state.save_request_data = None # Clear request
+             st.error(f"Received invalid save data format from client (expected list): {type(objects_to_save)}")
 
 
     except json.JSONDecodeError:
         st.error("Failed to decode save data from client. Data might be corrupted or empty.")
-        print("Received raw data:", save_data_from_js) # Log for debugging
-        st.session_state.save_request_data = None # Clear request
+        print("Received raw data:", save_data_from_js) # Log raw data
     except Exception as e:
         st.error(f"Error processing save: {e}")
         st.exception(e)
-        st.session_state.save_request_data = None # Clear request
+
+    # IMPORTANT: Clear the session state key regardless of success/failure
+    # to prevent reprocessing on the next rerun unless the button is clicked again.
+    st.session_state.js_save_processor = None
+
+    # Rerun AGAIN after processing save to reflect changes (new plot loaded, cache cleared etc.)
+    if save_processed_successfully:
+        st.rerun()
 
 
 # --- Main Area ---
 st.header("Shared 3D World")
-st.caption("Build side-by-side with others. Changes load on refresh/save.")
+st.caption("Build side-by-side with others. Saving adds a new plot to the right.")
 
 # --- Load and Prepare HTML ---
 html_file_path = 'index.html'
 html_content_with_state = None # Initialize
 
 try:
+    # --- Read the HTML template file ---
     with open(html_file_path, 'r', encoding='utf-8') as f:
         html_template = f.read()
 
-    # --- Inject Python state into JavaScript ---
+    # --- Prepare JavaScript code to inject state ---
+    # Ensure all state variables are correctly serialized as JSON
     js_injection_script = f"""
 <script>
     // Set global variables BEFORE the main script runs
     window.ALL_INITIAL_OBJECTS = {json.dumps(all_initial_objects)}; // All objects from all plots
     window.SELECTED_OBJECT_TYPE = {json.dumps(st.session_state.selected_object)};
     window.PLOT_WIDTH = {json.dumps(PLOT_WIDTH)};
-    window.NEXT_PLOT_X_OFFSET = {json.dumps(next_plot_x_offset)}; // Needed for save calculation
+    window.NEXT_PLOT_X_OFFSET = {json.dumps(next_plot_x_offset)}; // Needed for save calculation & ground size
+    // Basic logging to verify state in browser console
     console.log("Streamlit State Injected:", {{
         selectedObject: window.SELECTED_OBJECT_TYPE,
         initialObjectsCount: window.ALL_INITIAL_OBJECTS ? window.ALL_INITIAL_OBJECTS.length : 0,
@@ -279,18 +325,21 @@ try:
     }});
 </script>
 """
+    # --- Inject the script into the HTML template ---
+    # Replacing just before </head> is generally safe
     html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)
 
-    # --- Embed HTML Component ---
+    # --- Embed HTML Component (ONLY if HTML loading and preparation succeeded) ---
     components.html(
         html_content_with_state,
-        height=750,
+        height=750, # Adjust height as needed
         scrolling=False
     )
 
+# --- Error Handling ---
 except FileNotFoundError:
     st.error(f"CRITICAL ERROR: Could not find the file '{html_file_path}'.")
-    st.warning(f"Make sure `{html_file_path}` is in the same directory as `app.py` and `{SAVE_DIR}` exists.")
+    st.warning(f"Please make sure `{html_file_path}` is in the same directory as `app.py` and that the `{SAVE_DIR}` directory exists.")
 except Exception as e:
     st.error(f"An critical error occurred during HTML preparation or component rendering: {e}")
-    st.exception(e)
+    st.exception(e) # Show full traceback for debugging