2 files changed, 47 insertions, 10 deletions

diff --git a/train_cmd.c b/train_cmd.c
index 4539e83f2..b3273dc68 100644
--- a/train_cmd.c
+++ b/train_cmd.c
@@ -841,13 +841,15 @@ int32 CmdMoveRailVehicle(int x, int y, uint32 flags, uint32 p1, uint32 p2)
 		if (src_head->subtype == TS_Front_Engine)
 			UpdateTrainAcceleration(src_head);
 		InvalidateWindow(WC_VEHICLE_DETAILS, src_head->index);
+		InvalidateWindow(WC_VEHICLE_REFIT, src_head->index);
 
 		if (dst_head) {
 			if (dst_head->subtype == TS_Front_Engine)
 				UpdateTrainAcceleration(dst_head);
 			InvalidateWindow(WC_VEHICLE_DETAILS, dst_head->index);
-			/* Update the refit button */
+			/* Update the refit button and window */
 			InvalidateWindowWidget(WC_VEHICLE_VIEW, dst_head->index, 12);
+			InvalidateWindow(WC_VEHICLE_REFIT, dst_head->index);
 		}
 
 		/* I added this to so that the refit buttons get updated */
@@ -996,6 +998,7 @@ int32 CmdSellRailWagon(int x, int y, uint32 flags, uint32 p1, uint32 p2)
 				/* 5. If the train still exists, update its acceleration, window, etc. */
 				if (first != NULL && first->subtype == TS_Front_Engine) {
 					InvalidateWindow(WC_VEHICLE_DETAILS, first->index);
+					InvalidateWindow(WC_VEHICLE_REFIT, first->index);
 					UpdateTrainAcceleration(first);
 				}
 
@@ -1315,21 +1318,41 @@ int32 CmdRefitRailVehicle(int x, int y, uint32 flags, uint32 p1, uint32 p2)
 	cost = 0;
 	num = 0;
 
+	// newgrf stuff can change graphics when refitting
+	if (!(flags & DC_EXEC))
+		InvalidateWindow(WC_VEHICLE_DEPOT, v->tile);
+
 	do {
 		/* XXX: We also refit all the attached wagons en-masse if they
 		 * can be refitted. This is how TTDPatch does it.  TODO: Have
 		 * some nice [Refit] button near each wagon. --pasky */
 		if (!CanRefitTo(v, new_cid)) continue;
 
-		if (new_cid != v->cargo_type && v->cargo_cap != 0) {
-			cost += (_price.build_railvehicle >> 8);
-			num += v->cargo_cap;
+		if (v->cargo_cap != 0) {
+			RailVehicleInfo *rvi = RailVehInfo(v->engine_type);
+			uint16 amount = rvi->capacity;
+			CargoID old_cid = rvi->cargo_type;
+
+			/* the capacity depends on the cargo type, a rail vehicle
+			 * can carry twice as much mail/goods as normal cargo,
+			 * and four times as much passengers */
+			(old_cid == CT_PASSENGERS) ||
+			(amount <<= 1, old_cid == CT_MAIL || old_cid == CT_GOODS) ||
+			(amount <<= 1, true);
+			(new_cid == CT_PASSENGERS) ||
+			(amount >>= 1, new_cid == CT_MAIL || new_cid == CT_GOODS) ||
+			(amount >>= 1, true);
+
+			if (new_cid != v->cargo_type)
+				cost += (_price.build_railvehicle >> 8);
+			num += amount;
 			if (flags & DC_EXEC) {
 				//autorefitted train cars wants to keep the cargo
 				//it will be checked if the cargo is valid in CmdReplaceVehicle
 				if (!(SkipStoppedInDepotCheck))
 					v->cargo_count = 0;
 				v->cargo_type = new_cid;
+				v->cargo_cap = amount;
 				InvalidateWindow(WC_VEHICLE_DETAILS, v->index);
 			}
 		}
diff --git a/vehicle_gui.c b/vehicle_gui.c
index c83f287ec..ce67dc171 100644
--- a/vehicle_gui.c
+++ b/vehicle_gui.c
@@ -165,9 +165,9 @@ void DrawVehicleProfitButton(Vehicle *v, int x, int y)
 	DrawSprite(SPR_BLOT | ormod, x, y);
 }
 
-/** Draw the list of available refit options.
+/** Draw the list of available refit options for a consist.
  * Draw the list and highlight the selected refit option (if any)
- * @param *v vehicle(type) to get the refit-options of
+ * @param *v first vehicle in consist to get the refit-options of
  * @param sel selected refit cargo-type in the window
  * @return the cargo type that is hightlighted, CT_INVALID if none
  */
@@ -176,6 +176,7 @@ CargoID DrawVehicleRefitWindow(const Vehicle *v, int sel)
 	uint32 cmask;
 	CargoID cid, cargo = CT_INVALID;
 	int y = 25;
+	const Vehicle* u;
 #define show_cargo(ctype) { \
 		byte colour = 16; \
 		if (sel == 0) { \
@@ -187,10 +188,23 @@ CargoID DrawVehicleRefitWindow(const Vehicle *v, int sel)
 		y += 10; \
 }
 
-		/* Check if engine has custom refit or normal ones, and get its bitmasked value.
-		 * Now just and it with the bitmasked available cargo on the current landscape, and
-		* where the bits are set: those are available */
-		cmask = (_engine_refit_masks[v->engine_type] != 0) ? _engine_refit_masks[v->engine_type] : _default_refitmasks[v->type - VEH_Train];
+		/* Check if vehicle has custom refit or normal ones, and get its bitmasked value.
+		 * If its a train, 'or' this with the refit masks of the wagons. Now just 'and'
+		 * it with the bitmask of available cargo on the current landscape, and
+		 * where the bits are set: those are available */
+		cmask = 0;
+		u = v;
+		do {
+			if (_engine_refit_masks[u->engine_type] != 0) { // newgrf custom refit mask
+				cmask |= _engine_refit_masks[u->engine_type];
+			} else if (u->cargo_cap != 0) {
+				// rail wagons cant be refitted by default
+				if (v->type != VEH_Train || !(RailVehInfo(u->engine_type)->flags & RVI_WAGON))
+					cmask |= _default_refitmasks[v->type - VEH_Train];
+			}
+			u = u->next;
+		} while (v->type == VEH_Train && u != NULL);
+
 		cmask &= _landscape_global_cargo_mask[_opt_ptr->landscape];
 
 		/* Check which cargo has been selected from the refit window and draw list */