From fa4db3891f5d4b72b72341c4131f9067799a3011 Mon Sep 17 00:00:00 2001
From: Gerkiz <fs@nvfs.se>
Date: Mon, 11 Nov 2024 18:47:52 +0100
Subject: [PATCH] Mtn: fix locomotive issue

---
 maps/mountain_fortress_v3/functions.lua     |  2 +-
 maps/mountain_fortress_v3/icw/functions.lua | 16 +++++++++++++---
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/maps/mountain_fortress_v3/functions.lua b/maps/mountain_fortress_v3/functions.lua
index 4d8ca14a..55e4305b 100644
--- a/maps/mountain_fortress_v3/functions.lua
+++ b/maps/mountain_fortress_v3/functions.lua
@@ -2511,7 +2511,7 @@ function Public.on_player_changed_position(event)
         if not loco_surface or not loco_surface.valid then
             return
         end
-        if player.surface.index == loco_surface.index then
+        if player.character ~= nil and player.character.surface.index == loco_surface.index then
             local map_gen = loco_surface.map_gen_settings
             if player.position.y > map_gen.height then player.set_controller { type = 1, character = player.character } end
             if player.position.y < (-map_gen.height / 2) then player.set_controller { type = 1, character = player.character } end
diff --git a/maps/mountain_fortress_v3/icw/functions.lua b/maps/mountain_fortress_v3/icw/functions.lua
index 8d104bcc..1cfc14d4 100644
--- a/maps/mountain_fortress_v3/icw/functions.lua
+++ b/maps/mountain_fortress_v3/icw/functions.lua
@@ -972,7 +972,12 @@ function Public.reconstruct_all_trains(icw)
         if not validate_entity(wagon.entity) then
             icw.wagons[unit_number] = nil
             Public.request_reconstruction(icw)
-            goto cont
+            return
+        end
+
+        local locomotive = WPT.get('locomotive')
+        if not (locomotive and locomotive.valid) then
+            return
         end
 
         if not wagon.surface then
@@ -981,8 +986,13 @@ function Public.reconstruct_all_trains(icw)
         end
         local carriages = wagon.entity.train.carriages
 
-        Public.construct_train(icw, wagon.entity, carriages)
-        ::cont::
+        for _, carriage in pairs(carriages) do
+            if not icw.wagons[carriage.unit_number] then
+                Public.create_wagon(icw, carriage)
+            end
+        end
+
+        Public.construct_train(icw, locomotive, carriages)
     end
     delete_empty_surfaces(icw)
     return true