Update app.py

app.py

@@ -55,6 +55,8 @@ class FactoryGame:
                                    for _ in range(self.GRID_SIZE)]
         if "last_update" not in st.session_state:
             st.session_state.last_update = time.time()
+        if "selected_building" not in st.session_state:
+            st.session_state.selected_building = None
 
     def update_resources(self):
         current_time = time.time()
@@ -95,6 +97,13 @@ class FactoryGame:
         st.success(f"Placed {building.name}")
         return True
 
+def select_building(building_id: str):
+    st.session_state.selected_building = building_id
+
+def handle_cell_click(row: int, col: int, game: FactoryGame):
+    if st.session_state.selected_building:
+        game.place_building(row, col, st.session_state.selected_building)
+
 def main():
     st.set_page_config(page_title="Factory Game", layout="wide")
     
@@ -118,21 +127,28 @@ def main():
         
         # Building selection
         st.subheader("Buildings")
-        selected_building = None
         for building_id, building in game.buildings.items():
             if st.button(
                 f"{building.icon} {building.name}\n"
                 f"Cost: {', '.join(f'{cost} {res}' for res, cost in building.cost.items())}",
                 help=building.description,
-                key=f"building_{building_id}"
+                key=f"building_{building_id}",
+                type="secondary" if st.session_state.selected_building != building_id else "primary"
             ):
-                selected_building = building_id
+                select_building(building_id)
+
+        # Currently selected building
+        if st.session_state.selected_building:
+            st.caption(f"Selected: {game.buildings[st.session_state.selected_building].name}")
+        else:
+            st.caption("Select a building to place")
 
     with col1:
         # Game grid
         st.subheader("Factory Grid")
         
-        # Create a grid of buttons
+        # Create a grid of buttons using containers for better layout
+        grid_container = st.container()
         for row in range(game.GRID_SIZE):
             cols = st.columns(game.GRID_SIZE)
             for col, column in enumerate(cols):
@@ -141,14 +157,14 @@ def main():
                 
                 # Style the button based on what's there
                 label = building.icon if building else "⬜"
-                
-                if column.button(
-                    label,
-                    key=f"cell_{row}_{col}",
-                    help=building.name if building else "Empty space"
-                ):
-                    if selected_building:
-                        game.place_building(row, col, selected_building)
+                with column:
+                    if st.button(
+                        label,
+                        key=f"cell_{row}_{col}",
+                        help=building.name if building else "Empty space",
+                        use_container_width=True
+                    ):
+                        handle_cell_click(row, col, game)
     
     # Update resources periodically
     game.update_resources()