Update app.py

app.py

@@ -31,15 +31,23 @@ def load_world_map():
 
 def save_world_map(world_data):
     """Saves the world map metadata."""
-    with open(WORLD_MAP_FILE, 'w') as f:
-        json.dump(world_data, f, indent=4)
+    try:
+        with open(WORLD_MAP_FILE, 'w') as f:
+            json.dump(world_data, f, indent=4)
+    except Exception as e:
+        st.error(f"Error saving world map: {e}")
+
 
 def save_space_data(space_id, objects_data):
     """Saves the object data for a specific space."""
     file_path = os.path.join(SAVE_DIR, f"{space_id}.json")
-    with open(file_path, 'w') as f:
-        # Store objects directly, could add metadata later
-        json.dump({"objects": objects_data}, f, indent=4)
+    try:
+        with open(file_path, 'w') as f:
+            # Store objects directly, could add metadata later
+            json.dump({"objects": objects_data}, f, indent=4)
+    except Exception as e:
+        st.error(f"Error saving space data for {space_id}: {e}")
+
 
 def load_space_data(space_id):
     """Loads object data for a specific space."""
@@ -52,6 +60,9 @@ def load_space_data(space_id):
         except json.JSONDecodeError:
             st.error(f"Error reading space file {space_id}.json.")
             return []
+        except Exception as e:
+            st.error(f"Error loading space data for {space_id}: {e}")
+            return []
     return [] # Return empty list if file doesn't exist
 
 def find_next_available_grid_slot(world_data):
@@ -59,18 +70,38 @@ def find_next_available_grid_slot(world_data):
     occupied = set((d["grid_x"], d["grid_y"]) for d in world_data.get("spaces", {}).values())
     x, y = 0, 0
     dx, dy = 0, -1
-    steps = 0
-    limit = 1
-    count = 0
-    while (x, y) in occupied:
-        if x == y or (x < 0 and x == -y) or (x > 0 and x == 1-y):
-             dx, dy = -dy, dx # Change direction (spiral)
+    limit = 1 # How many steps in current direction
+    steps = 0 # Steps taken in current direction
+    leg = 0   # Current leg of the spiral (0=right, 1=up, 2=left, 3=down)
+
+    while True:
+        if (x, y) not in occupied:
+            return x, y
+
+        # Move
         x, y = x + dx, y + dy
-        count += 1
-        if count > 1000: # Safety break
-             st.error("Could not find empty grid slot easily!")
+        steps += 1
+
+        # Check if direction needs changing
+        if steps == limit:
+            steps = 0
+            leg = (leg + 1) % 4
+            if leg == 0: # Right
+                dx, dy = 1, 0
+            elif leg == 1: # Up
+                dx, dy = 0, 1
+                limit += 1 # Increase steps after moving up
+            elif leg == 2: # Left
+                dx, dy = -1, 0
+            elif leg == 3: # Down
+                dx, dy = 0, -1
+                limit += 1 # Increase steps after moving down
+
+        # Safety break (increase if expecting huge maps)
+        if limit > 100:
+             st.error("Could not find empty grid slot easily! Map too full?")
              return None, None
-    return x, y
+
 
 # --- Minimap Generation ---
 def generate_minimap(world_data, current_space_id=None):
@@ -78,33 +109,44 @@ def generate_minimap(world_data, current_space_id=None):
     if not spaces:
         return None # No map if no spaces saved
 
-    coords = [(d["grid_x"], d["grid_y"]) for d in spaces.values()]
-    min_x = min(c[0] for c in coords)
-    max_x = max(c[0] for c in coords)
-    min_y = min(c[1] for c in coords)
-    max_y = max(c[1] for c in coords)
-
-    img_width = (max_x - min_x + 1) * MINIMAP_CELL_SIZE
-    img_height = (max_y - min_y + 1) * MINIMAP_CELL_SIZE
-
-    img = Image.new('RGB', (img_width, img_height), color = 'lightgrey')
-    draw = ImageDraw.Draw(img)
-
-    for space_id, data in spaces.items():
-        cell_x = (data["grid_x"] - min_x) * MINIMAP_CELL_SIZE
-        cell_y = (data["grid_y"] - min_y) * MINIMAP_CELL_SIZE
-        color = "blue"
-        if space_id == current_space_id:
-            color = "red" # Highlight current space
-        elif current_space_id is None and space_id == list(spaces.keys())[0]: # Highlight first if none selected
-             color = "red"
-
-        draw.rectangle(
-            [cell_x, cell_y, cell_x + MINIMAP_CELL_SIZE -1, cell_y + MINIMAP_CELL_SIZE -1],
-            fill=color, outline="black"
-        )
-
-    return img
+    try:
+        coords = [(d["grid_x"], d["grid_y"]) for d in spaces.values()]
+        # Handle case where coords might be empty if spaces dict is malformed
+        if not coords: return None
+
+        min_x = min(c[0] for c in coords)
+        max_x = max(c[0] for c in coords)
+        min_y = min(c[1] for c in coords)
+        max_y = max(c[1] for c in coords)
+
+        # Add padding around the edges
+        padding = 1
+        img_width = (max_x - min_x + 1 + 2 * padding) * MINIMAP_CELL_SIZE
+        img_height = (max_y - min_y + 1 + 2 * padding) * MINIMAP_CELL_SIZE
+
+        img = Image.new('RGB', (img_width, img_height), color = 'lightgrey')
+        draw = ImageDraw.Draw(img)
+
+        for space_id, data in spaces.items():
+            # Calculate position including padding offset
+            cell_x = (data["grid_x"] - min_x + padding) * MINIMAP_CELL_SIZE
+            cell_y = (data["grid_y"] - min_y + padding) * MINIMAP_CELL_SIZE
+            color = "blue"
+            if space_id == current_space_id:
+                color = "red" # Highlight current space
+
+            draw.rectangle(
+                [cell_x, cell_y, cell_x + MINIMAP_CELL_SIZE -1, cell_y + MINIMAP_CELL_SIZE -1],
+                fill=color, outline="black"
+            )
+            # Optional: Draw space name/ID lightly
+            # draw.text((cell_x + 1, cell_y + 1), data.get("name", sid[:2]), fill="white", font_size=8)
+
+
+        return img
+    except Exception as e:
+        st.error(f"Error generating minimap: {e}")
+        return None
 
 
 # --- Page Config ---
@@ -122,6 +164,9 @@ if 'space_name' not in st.session_state:
      st.session_state.space_name = ""
 if 'initial_objects' not in st.session_state:
      st.session_state.initial_objects = [] # Objects to load into JS
+if 'trigger_rerun' not in st.session_state:
+     st.session_state.trigger_rerun = 0 # Counter to force reruns cleanly
+
 
 # --- Load initial world data ---
 world_data = load_world_map()
@@ -130,50 +175,83 @@ world_data = load_world_map()
 query_params = st.query_params.to_dict()
 save_data_encoded = query_params.get("save_data")
 
-save_triggered = False
+
 if save_data_encoded:
+    st.session_state.trigger_rerun += 1 # Increment to signal change
+    save_successful = False
     try:
         save_data_json = urllib.parse.unquote(save_data_encoded[0]) # Get first value if list
         objects_to_save = json.loads(save_data_json)
 
-        space_id_to_save = query_params.get("space_id", [st.session_state.current_space_id])[0] # Get from URL or state
-        space_name_to_save = query_params.get("space_name", [st.session_state.space_name])[0]
+        # Get ID/Name from URL or fall back to session state if not in URL (e.g., if redirect failed partially)
+        space_id_to_save = query_params.get("space_id", [st.session_state.current_space_id])[0]
+        space_name_to_save = urllib.parse.unquote(query_params.get("space_name", [st.session_state.space_name])[0])
+
 
-        if not space_id_to_save:
+        if not space_id_to_save or space_id_to_save == 'null': # JS might send 'null' string
             space_id_to_save = str(uuid.uuid4()) # Create new ID
+            is_new_space = True
             st.session_state.current_space_id = space_id_to_save # Update state
+        else:
+            is_new_space = False
+
+        # Save the actual object data
+        save_space_data(space_id_to_save, objects_to_save)
+
+        # Update world map only if needed (new space or name change)
+        map_updated = False
+        world_spaces = world_data.setdefault("spaces", {})
+
+        if is_new_space:
             grid_x, grid_y = find_next_available_grid_slot(world_data)
             if grid_x is not None:
-                 world_data.setdefault("spaces", {})[space_id_to_save] = {
+                 world_spaces[space_id_to_save] = {
                      "grid_x": grid_x,
                      "grid_y": grid_y,
-                     "name": space_name_to_save or f"Space {len(world_data.get('spaces',{}))+1}"
+                     "name": space_name_to_save or f"Space_{space_id_to_save[:6]}"
                  }
-                 save_world_map(world_data)
+                 map_updated = True
             else:
                  st.error("Failed to assign grid position!")
+        elif space_id_to_save in world_spaces and world_spaces[space_id_to_save].get("name") != space_name_to_save:
+             # Update name if it changed for existing space
+             world_spaces[space_id_to_save]["name"] = space_name_to_save
+             map_updated = True
 
+        if map_updated:
+             save_world_map(world_data)
 
-        # Save the actual object data
-        save_space_data(space_id_to_save, objects_to_save)
-        st.success(f"Space '{space_name_to_save or space_id_to_save}' saved successfully!")
+        st.session_state.space_name = space_name_to_save # Ensure state matches saved name
 
-        # Update name in world map if it changed and space exists
-        if space_id_to_save in world_data.get("spaces", {}) and space_name_to_save:
-             world_data["spaces"][space_id_to_save]["name"] = space_name_to_save
-             save_world_map(world_data)
+        st.success(f"Space '{space_name_to_save or space_id_to_save}' saved successfully!")
+        save_successful = True
 
-        # IMPORTANT: Clear query param to prevent resave on refresh
-        st.query_params.clear()
-        save_triggered = True # Flag to maybe skip immediate rerun if needed below
 
     except Exception as e:
         st.error(f"Error processing save data: {e}")
-        st.exception(e) # Show traceback
+        st.exception(e) # Show full traceback for debugging
+
+    # IMPORTANT: Clear query param to prevent resave on refresh/rerun
+    st.query_params.clear()
+
+    # Force a rerun IF save was processed, to reload objects and clear URL state visually
+    # Do this AFTER clearing query_params
+    if save_successful:
+        # No need to explicitly call st.rerun() here, Streamlit's flow after
+        # clearing query params and updating state often handles it.
+        # If visual updates lag, uncommenting st.rerun() might be necessary,
+        # but can sometimes cause double reruns.
+        # st.rerun()
+        pass
 
-# Need to reload objects if just saved or loading a new space
-if 'current_space_id' in st.session_state and st.session_state.current_space_id:
+
+# --- Load Space Data if ID is set ---
+# This runs on every script execution, including after save/load actions
+if st.session_state.current_space_id:
      st.session_state.initial_objects = load_space_data(st.session_state.current_space_id)
+     # Ensure name is loaded if space ID exists but name state is empty
+     if not st.session_state.space_name and st.session_state.current_space_id in world_data.get("spaces", {}):
+          st.session_state.space_name = world_data["spaces"][st.session_state.current_space_id].get("name", "")
 
 
 # --- Sidebar Controls ---
@@ -184,62 +262,86 @@ with st.sidebar:
     st.subheader("Manage Spaces")
     saved_spaces = list(world_data.get("spaces", {}).items()) # List of (id, data) tuples
     space_options = {sid: data.get("name", f"Unnamed ({sid[:6]}...)") for sid, data in saved_spaces}
-    space_options["_new_"] = "✨ Create New Space ✨" # Special option
+    options_list = ["_new_"] + list(space_options.keys())
+
+    # Determine current index for selectbox
+    current_selection_index = 0 # Default to "Create New"
+    if st.session_state.current_space_id in space_options:
+        try:
+            current_selection_index = options_list.index(st.session_state.current_space_id)
+        except ValueError:
+            current_selection_index = 0 # Fallback if ID somehow not in list
 
-    selected_space_display = st.selectbox(
+    selected_space_option = st.selectbox(
         "Load or Create Space:",
-        options = ["_new_"] + list(space_options.keys()),
+        options = options_list,
         format_func = lambda x: space_options.get(x, "Select...") if x != "_new_" else "✨ Create New Space ✨",
-        index=0, # Default to Create New
-        key="space_selection_key" # Unique key might be needed if dynamically changing options
-        # Note: Changing this will trigger rerun. Need logic below to handle load.
+        index=current_selection_index,
+        key="space_selection_widget" # Use a dedicated key for the widget
     )
 
-    # Handle Load/Create based on selection
-    if st.session_state.space_selection_key != "_new_":
-        # Load existing selected
-        if st.session_state.current_space_id != st.session_state.space_selection_key:
-             st.session_state.current_space_id = st.session_state.space_selection_key
-             st.session_state.initial_objects = load_space_data(st.session_state.current_space_id)
-             st.session_state.space_name = world_data["spaces"][st.session_state.current_space_id].get("name", "")
-             st.rerun() # Rerun to load data and update JS injection
-    elif st.session_state.space_selection_key == "_new_" and st.session_state.current_space_id is not None:
-         # Handle switch from existing to "Create New"
-         st.session_state.current_space_id = None
-         st.session_state.initial_objects = []
-         st.session_state.space_name = ""
-         st.rerun()
-
-
+    # --- Handle Load/Create Logic ---
+    # Compare widget state to session state to detect user change
+    if selected_space_option != st.session_state.get('_last_selected_space', None):
+        st.session_state._last_selected_space = selected_space_option # Track selection change
+        if selected_space_option == "_new_":
+            # User explicitly selected "Create New"
+            if st.session_state.current_space_id is not None: # Check if switching *from* an existing space
+                st.session_state.current_space_id = None
+                st.session_state.initial_objects = []
+                st.session_state.space_name = ""
+                st.rerun() # Rerun to clear the space
+        else:
+            # User selected an existing space
+            if st.session_state.current_space_id != selected_space_option:
+                st.session_state.current_space_id = selected_space_option
+                # Loading data and name happens naturally at top of script now
+                st.rerun() # Rerun to reflect the load
+
+    # --- Name Input ---
     current_name = st.text_input(
          "Current Space Name:",
-         value=st.session_state.space_name,
+         value=st.session_state.space_name, # Reflect current loaded/new name state
          key="current_space_name_input"
     )
-    # Update state immediately if name changes
+    # Update session state if user types a new name
     if current_name != st.session_state.space_name:
         st.session_state.space_name = current_name
-        # If editing an existing space, maybe update world_map.json immediately or on save? Let's do on save for simplicity.
-
+        # Note: The name is only saved to world_map.json when the JS save is triggered
 
     st.info(f"Current Space ID: {st.session_state.current_space_id or 'None (New)'}")
-    st.caption("Saving uses the name above. A unique ID is assigned automatically for new spaces.")
+    st.caption("Click 'Save Work' in the 3D view to save changes.")
     st.markdown("---")
 
 
     # --- Object Placement Controls ---
     st.subheader("Place Objects")
     object_types = ["None", "Simple House", "Tree", "Rock", "Fence Post"]
-    selected = st.selectbox(
+    # Ensure selectbox reflects current state
+    current_object_index = 0
+    try:
+        current_object_index = object_types.index(st.session_state.selected_object)
+    except ValueError:
+        st.session_state.selected_object = "None" # Reset if invalid
+
+    selected_object_type_widget = st.selectbox(
         "Select Object to Place:",
         options=object_types,
-        key='selected_object'
+        index=current_object_index,
+        key="selected_object_widget"
     )
+    # Update state based on widget interaction
+    if selected_object_type_widget != st.session_state.selected_object:
+         st.session_state.selected_object = selected_object_type_widget
+         st.rerun() # Rerun to update JS injection
+
     st.markdown("---")
 
     # --- Minimap ---
     st.subheader("World Minimap")
-    minimap_img = generate_minimap(world_data, st.session_state.current_space_id)
+    # Regenerate map data fresh each time
+    current_world_data_for_map = load_world_map()
+    minimap_img = generate_minimap(current_world_data_for_map, st.session_state.current_space_id)
     if minimap_img:
         st.image(minimap_img, caption="Blue: Saved Spaces, Red: Current", use_column_width=True)
     else:
@@ -252,12 +354,15 @@ st.caption(f"Editing: {st.session_state.space_name or 'New Space'}")
 
 # --- Load and Prepare HTML ---
 html_file_path = 'index.html'
+html_content_with_state = None # Initialize to None
 
 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
@@ -265,28 +370,31 @@ try:
     window.INITIAL_OBJECTS = {json.dumps(st.session_state.initial_objects)};
     window.CURRENT_SPACE_ID = {json.dumps(st.session_state.current_space_id)};
     window.CURRENT_SPACE_NAME = {json.dumps(st.session_state.space_name)};
-    console.log("Streamlit State:", {{
+    // Basic logging to verify state in browser console
+    console.log("Streamlit State Injected:", {{
         selectedObject: window.SELECTED_OBJECT_TYPE,
-        initialObjects: window.INITIAL_OBJECTS,
+        initialObjectsCount: window.INITIAL_OBJECTS ? window.INITIAL_OBJECTS.length : 0,
         spaceId: window.CURRENT_SPACE_ID,
         spaceName: window.CURRENT_SPACE_NAME
     }});
 </script>
 """
-    # Insert the injection script just before the closing </head>
-    # Using placeholder is safer if you modify index.html: html_content_with_state = html_template.replace('</head>', js_injection_script + '\n</head>', 1)
-
+    # --- 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, # Adjust height as needed
         scrolling=False
     )
 
+# --- Error Handling ---
 except FileNotFoundError:
-    st.error(f"Error: Could not find the file '{html_file_path}'.")
+    st.error(f"CRITICAL ERROR: Could not find the file '{html_file_path}'.")
     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 error occurred: {e}")
-    st.exception(e)
+    st.error(f"An critical error occurred during HTML preparation or component rendering: {e}")
+    st.exception(e) # Show full traceback for debugging
+    # Do NOT attempt to render components.html if html_content_with_state is not defined