From 3d7103ce3153f29173d8c811b36e664d5f7792a3 Mon Sep 17 00:00:00 2001
From: Filip Lange <filip.lange@gmail.com>
Date: Fri, 22 Sep 2023 22:52:20 +0200
Subject: [PATCH] Add destroy train button

---
 tests/test_game.py   | 17 +++++++++++++++++
 trainfinity2/box.py  |  3 +++
 trainfinity2/game.py | 36 +++++++++++++++++++++---------------
 trainfinity2/gui.py  | 11 ++++++-----
 4 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/tests/test_game.py b/tests/test_game.py
index b43387a..79cde1e 100644
--- a/tests/test_game.py
+++ b/tests/test_game.py
@@ -1265,3 +1265,20 @@ def test_success(self, game: Game):
         assert len(train.wagons) == 1
         game._create_wagon_for_selected_train()
         assert len(train.wagons) == 2
+
+
+def test_destroy_train_button(game: Game):
+    create_objects(
+        game.grid,
+        """
+        . M . F .
+
+        .-S-.-S-.
+        """,
+    )
+    train = game._create_train(*game.grid.station_from_position.values())
+    train.selected = True
+
+    assert len(game.trains) == 1
+    game.gui.boxes["DESTROY\nTRAIN"].click()
+    assert len(game.trains) == 0
diff --git a/trainfinity2/box.py b/trainfinity2/box.py
index 2f2b1cc..6df8584 100644
--- a/trainfinity2/box.py
+++ b/trainfinity2/box.py
@@ -11,3 +11,6 @@ class Box:
     callback: Callable[..., None]
     callback_args: list
     mode: Mode | None = None
+
+    def click(self):
+        self.callback(*self.callback_args)
diff --git a/trainfinity2/game.py b/trainfinity2/game.py
index 8dd02af..01a8242 100644
--- a/trainfinity2/game.py
+++ b/trainfinity2/game.py
@@ -81,6 +81,7 @@ def setup(self, terrain: Terrain):
             Box("SIGNAL", self._set_mode, [Mode.SIGNAL], Mode.SIGNAL),
             Box("DESTROY", self._set_mode, [Mode.DESTROY], Mode.DESTROY),
             Box("+WAGON", self._create_wagon_for_selected_train, []),
+            Box("DESTROY\nTRAIN", self._destroy_selected_train, []),
         ]
 
         self.gui_camera = Camera()
@@ -245,23 +246,28 @@ def _create_train(self, station1: Station, station2: Station):
         train.selected = True
         return train
 
+    @property
+    def _selected_train(self) -> Train | None:
+        return next((train for train in self.trains if train.selected), None)
+
+    def _destroy_selected_train(self):
+        if train := self._selected_train:
+            self._destroy_train(train)
+
     def _create_wagon_for_selected_train(self):
-        for train in self.trains:
-            if train.selected:
-                station1 = self.grid.station_from_position[train.first_station_position]
-                station2 = self.grid.station_from_position[
-                    train.second_station_position
-                ]
-                shortest_station_length = min(
-                    len(station1.positions),
-                    len(station2.positions),
+        if train := self._selected_train:
+            station1 = self.grid.station_from_position[train.first_station_position]
+            station2 = self.grid.station_from_position[train.second_station_position]
+            shortest_station_length = min(
+                len(station1.positions),
+                len(station2.positions),
+            )
+            if shortest_station_length > len(train.wagons) + 1:
+                train.add_wagon()
+            else:
+                self.gui.toast(
+                    f"Wagon count ({len(train.wagons)}) must be less than shortest station length ({shortest_station_length})"
                 )
-                if shortest_station_length > len(train.wagons) + 1:
-                    train.add_wagon()
-                else:
-                    self.gui.toast(
-                        f"Wagon count ({len(train.wagons)}) must be less than shortest station length ({shortest_station_length})"
-                    )
 
     def on_mouse_motion(self, x: int, y: int, dx: int, dy: int):
         events: list[Event] = []
diff --git a/trainfinity2/gui.py b/trainfinity2/gui.py
index b231c25..a3e86de 100644
--- a/trainfinity2/gui.py
+++ b/trainfinity2/gui.py
@@ -17,7 +17,7 @@
 class Gui:
     def __init__(self, camera: Camera, boxes: list[Box]) -> None:
         self.camera = camera
-        self._boxes = boxes
+        self.boxes = {box.text: box for box in boxes}
         self.mouse_press_box: Box | None = None
         self._enabled = True
         self._shape_element_list: arcade.ShapeElementList = arcade.ShapeElementList()
@@ -74,7 +74,7 @@ def refresh(self):
             changing color of the boxes when switching active mode."""
             self._shape_element_list = arcade.ShapeElementList()
             self._sprite_list = arcade.SpriteList()
-            for i, box in enumerate(self._boxes):
+            for i, box in enumerate(self.boxes.values()):
                 self._create_box(box, i)
 
             self.refresh_text()
@@ -95,7 +95,7 @@ def _is_inside(self, x, y, index):
     def on_mouse_press(self, x, y) -> bool:
         if self._enabled:
             with self.camera:
-                for i, box in enumerate(self._boxes):
+                for i, box in enumerate(self.boxes.values()):
                     if self._is_inside(x, y, i):
                         self.mouse_press_box = box
                         self.refresh()
@@ -105,10 +105,10 @@ def on_mouse_press(self, x, y) -> bool:
     def on_mouse_release(self, x, y) -> None:
         if self._enabled:
             with self.camera:
-                for i, box in enumerate(self._boxes):
+                for i, box in enumerate(self.boxes.values()):
                     if self._is_inside(x, y, i):
                         if self.mouse_press_box == box:
-                            box.callback(*box.callback_args)
+                            box.click()
         self.mouse_press_box = None
         self.refresh()
 
@@ -143,6 +143,7 @@ def _create_box(self, box: Box, index: int):
                 color=BOX_TEXT_COLOR,
                 width=int(BOX_SIZE_PIXELS),
                 align="center",
+                anchor_y="center",
             )
         )