From 541572e4150b91199425425fa5677c4e9940cb5e Mon Sep 17 00:00:00 2001
From: Spythere <spythere@protonmail.com>
Date: Mon, 15 Apr 2024 18:32:19 +0200
Subject: [PATCH] refactor: stock & vehicle typings

---
 src/components/app/MainContainer.vue          |  4 +-
 src/components/sections/InputsSection.vue     | 12 ++--
 src/components/sections/TrainImageSection.vue | 44 +++++++++----
 src/components/tabs/StockGeneratorTab.vue     | 11 +++-
 src/components/tabs/StockListTab.vue          | 65 ++++++++++---------
 src/components/tabs/WikiListTab.vue           | 37 +++++++----
 src/components/utils/StockThumbnails.vue      | 25 ++++---
 src/mixins/stockMixin.ts                      | 29 ++++-----
 src/mixins/stockPreviewMixin.ts               | 21 +++---
 src/store.ts                                  | 15 +++--
 src/styles/global.scss                        |  2 +-
 src/types.ts                                  | 16 ++---
 src/utils/vehicleUtils.ts                     | 27 ++++----
 13 files changed, 175 insertions(+), 133 deletions(-)

diff --git a/src/components/app/MainContainer.vue b/src/components/app/MainContainer.vue
index a5ab0d5..6b439cf 100644
--- a/src/components/app/MainContainer.vue
+++ b/src/components/app/MainContainer.vue
@@ -42,7 +42,9 @@ main {
 
 @media screen and (max-width: $breakpointMd) {
   main {
-    display: block;
+    display: flex;
+    flex-direction: column;
+    gap: 1em;
   }
 }
 </style>
diff --git a/src/components/sections/InputsSection.vue b/src/components/sections/InputsSection.vue
index a7728d7..5d7d6b8 100644
--- a/src/components/sections/InputsSection.vue
+++ b/src/components/sections/InputsSection.vue
@@ -28,7 +28,8 @@
             {{ $t('inputs.input-vehicle') }}
           </option>
           <option v-for="loco in locoOptions" :value="loco" :key="loco.type">
-            {{ loco.type }}<b v-if="loco.restrictions['sponsorOnly']">*</b>
+            {{ loco.type
+            }}<b v-if="loco.sponsorOnlyTimestamp && loco.sponsorOnlyTimestamp > Date.now()">*</b>
           </option>
         </select>
       </div>
@@ -59,7 +60,8 @@
           </option>
 
           <option v-for="car in carOptions" :value="car" :key="car.type">
-            {{ car.type }}<b v-if="car.restrictions['sponsorOnly']">*</b>
+            {{ car.type
+            }}<b v-if="car.sponsorOnlyTimestamp && car.sponsorOnlyTimestamp > Date.now()">*</b>
           </option>
         </select>
       </div>
@@ -84,6 +86,7 @@
           <option :value="null" v-if="!store.chosenCar || !store.chosenCar.loadable">
             {{ $t('inputs.no-cargo-available') }}
           </option>
+
           <option :value="null" v-else>{{ $t('inputs.cargo-empty') }}</option>
 
           <option v-for="cargo in store.chosenCar?.cargoTypes" :value="cargo" :key="cargo.id">
@@ -193,10 +196,7 @@ export default defineComponent({
     removeVehicle() {
       if (this.store.stockList.length == 0) return;
 
-      const lastStock = this.store.stockList.slice(-1)[0];
-
-      if (lastStock.count > 1) lastStock.count--;
-      else this.store.stockList.splice(-1);
+      this.store.stockList.splice(-1);
     },
 
     switchVehicles() {
diff --git a/src/components/sections/TrainImageSection.vue b/src/components/sections/TrainImageSection.vue
index 39dd75a..0538168 100644
--- a/src/components/sections/TrainImageSection.vue
+++ b/src/components/sections/TrainImageSection.vue
@@ -4,8 +4,11 @@
       <img
         :src="getThumbnailURL(store.chosenVehicle.type, 'small')"
         :data-preview-active="store.chosenVehicle !== null"
-        :data-sponsor-only="store.chosenVehicle?.restrictions.sponsorOnly"
-        :data-team-only="store.chosenVehicle?.restrictions.teamOnly"
+        :data-sponsor-only="
+          store.chosenVehicle.sponsorOnlyTimestamp &&
+          store.chosenVehicle.sponsorOnlyTimestamp > Date.now()
+        "
+        :data-team-only="store.chosenVehicle.teamOnly"
         @click="onImageClick"
         @keydown.enter="onImageClick"
         @error="onImageError"
@@ -41,17 +44,23 @@
             }}
           </div>
 
-          <b v-if="store.chosenVehicle.restrictions.sponsorOnly > 0" class="sponsor-only">{{
-            $t('preview.sponsor-only', [
-              new Date(store.chosenVehicle.restrictions['sponsorOnly']).toLocaleDateString(
-                $i18n.locale == 'pl' ? 'pl-PL' : 'en-GB'
-              ),
-            ])
-          }}</b>
+          <b
+            v-if="
+              store.chosenVehicle.sponsorOnlyTimestamp &&
+              store.chosenVehicle.sponsorOnlyTimestamp > Date.now()
+            "
+            class="sponsor-only"
+          >
+            {{
+              $t('preview.sponsor-only', [
+                new Date(store.chosenVehicle.sponsorOnlyTimestamp).toLocaleDateString(
+                  $i18n.locale == 'pl' ? 'pl-PL' : 'en-GB'
+                ),
+              ])
+            }}
+          </b>
 
-          <b v-if="store.chosenVehicle.restrictions['teamOnly']" class="team-only">{{
-            $t('preview.team-only')
-          }}</b>
+          <b v-if="store.chosenVehicle.teamOnly" class="team-only">{{ $t('preview.team-only') }}</b>
         </div>
       </div>
     </div>
@@ -120,8 +129,7 @@ export default defineComponent({
   align-items: center;
   text-align: center;
 
-  margin-top: 1em;
-  height: 250px;
+  min-height: 250px;
 
   & > div {
     max-width: 100%;
@@ -152,6 +160,14 @@ img {
   background-color: $bgColor;
 }
 
+.sponsor-only {
+  color: $sponsorColor;
+}
+
+.team-only {
+  color: $teamColor;
+}
+
 .image-info {
   font-size: 1.1em;
   padding: 0.5em;
diff --git a/src/components/tabs/StockGeneratorTab.vue b/src/components/tabs/StockGeneratorTab.vue
index 30aa139..2ec54c8 100644
--- a/src/components/tabs/StockGeneratorTab.vue
+++ b/src/components/tabs/StockGeneratorTab.vue
@@ -119,6 +119,7 @@ import { useStore } from '../../store';
 
 import stockMixin from '../../mixins/stockMixin';
 import { ICargo, ICarWagon, IStock } from '../../types';
+import { isTractionUnit } from '../../utils/vehicleUtils';
 
 export default defineComponent({
   name: 'stock-generator',
@@ -186,8 +187,8 @@ export default defineComponent({
       if (!this.isCarGroupingEnabled) return false;
 
       stockList.sort((s1, s2) => {
-        return (s1.constructionType + s1.cargo?.id).localeCompare(
-          s2.constructionType + s2.cargo?.id
+        return (s1.vehicleRef.constructionType + s1.cargo?.id).localeCompare(
+          s2.vehicleRef.constructionType + s2.cargo?.id
         );
       });
     },
@@ -238,7 +239,11 @@ export default defineComponent({
       };
 
       for (let i = 0; i < 10; i++) {
-        this.store.stockList.splice(this.store.stockList[0]?.isLoco ? 1 : 0);
+        this.store.stockList.splice(
+          this.store.stockList.length > 0 && isTractionUnit(this.store.stockList[0].vehicleRef)
+            ? 1
+            : 0
+        );
 
         let carCount = 0;
         const maxWeight =
diff --git a/src/components/tabs/StockListTab.vue b/src/components/tabs/StockListTab.vue
index deffaf0..24d46ef 100644
--- a/src/components/tabs/StockListTab.vue
+++ b/src/components/tabs/StockListTab.vue
@@ -132,7 +132,7 @@
         (!)
         {{
           $t('stocklist.warning-team-only-vehicle', [
-            teamOnlyVehicles.map((v) => v.type).join(', '),
+            teamOnlyVehicles.map((v) => v.vehicleRef.type).join(', '),
           ])
         }}
       </div>
@@ -169,7 +169,7 @@
           <li
             v-for="(stock, i) in store.stockList"
             :key="stock.id"
-            :class="{ loco: stock.isLoco }"
+            :class="{ loco: isTractionUnit(stock.vehicleRef) }"
             tabindex="0"
             @click="onListItemClick(i)"
             @keydown.enter="onListItemClick(i)"
@@ -192,22 +192,29 @@
 
               <span
                 class="stock-info-type"
-                :data-sponsor-only="stock.restrictions.sponsorOnly"
-                :data-team-only="stock.restrictions.teamOnly"
+                :data-sponsor-only="
+                  stock.vehicleRef.sponsorOnlyTimestamp &&
+                  stock.vehicleRef.sponsorOnlyTimestamp > Date.now()
+                "
+                :data-team-only="stock.vehicleRef.teamOnly"
               >
-                {{ stock.isLoco ? stock.type : getCarSpecFromType(stock.type) }}
+                {{
+                  isTractionUnit(stock.vehicleRef)
+                    ? stock.vehicleRef.type
+                    : getCarSpecFromType(stock.vehicleRef.type)
+                }}
               </span>
 
               <span class="stock-info-cargo" v-if="stock.cargo">
                 {{ stock.cargo.id }}
               </span>
 
-              <span class="stock-info-length">{{ stock.length }}m</span>
+              <span class="stock-info-length">{{ stock.vehicleRef.length }}m</span>
 
               <span class="stock-info-mass">
-                {{ ((stock.weight + (stock.cargo?.weight ?? 0)) / 1000).toFixed(1) }}t
+                {{ ((stock.vehicleRef.weight + (stock.cargo?.weight ?? 0)) / 1000).toFixed(1) }}t
               </span>
-              <span class="stock-info-speed">{{ stock.maxSpeed }}km/h</span>
+              <span class="stock-info-speed">{{ stock.vehicleRef.maxSpeed }}km/h</span>
             </div>
           </li>
         </transition-group>
@@ -226,6 +233,7 @@ import stockPreviewMixin from '../../mixins/stockPreviewMixin';
 import StockThumbnails from '../utils/StockThumbnails.vue';
 import stockMixin from '../../mixins/stockMixin';
 import Checkbox from '../common/Checkbox.vue';
+import { isTractionUnit } from '../../utils/vehicleUtils';
 
 export default defineComponent({
   name: 'stock-list',
@@ -250,7 +258,7 @@ export default defineComponent({
 
   computed: {
     chosenRealComposition() {
-      const currentStockString = this.store.stockList.map((s) => s.type).join(';');
+      const currentStockString = this.store.stockList.map((s) => s.vehicleRef.type).join(';');
 
       return this.store.realCompositionList.find((rc) => rc.stockString == currentStockString);
     },
@@ -265,7 +273,9 @@ export default defineComponent({
       return this.store.stockList
         .map((stock, i) => {
           let stockTypeStr =
-            stock.isLoco || !stock.cargo ? stock.type : `${stock.type}:${stock.cargo.id}`;
+            isTractionUnit(stock.vehicleRef) || !stock.cargo
+              ? stock.vehicleRef.type
+              : `${stock.vehicleRef.type}:${stock.cargo.id}`;
 
           if (i == 0 && (includeColdStart || includeDoubleManned))
             return `${stockTypeStr},${includeColdStart ? 'c' : ''}${includeDoubleManned ? 'd' : ''}`;
@@ -300,22 +310,24 @@ export default defineComponent({
       return (
         !this.store.isTrainPassenger &&
         this.store.stockList.length > 1 &&
-        !this.store.stockList.every((stock) => stock.isLoco) &&
-        this.store.stockList.some((stock) => stock.isLoco && stock.type.startsWith('EP'))
+        !this.store.stockList.every((stock) => isTractionUnit(stock.vehicleRef)) &&
+        this.store.stockList.some(
+          (stock) => isTractionUnit(stock.vehicleRef) && stock.vehicleRef.type.startsWith('EP')
+        )
       );
     },
 
     locoCountExceeded() {
       return (
         this.store.stockList.reduce((acc, stock) => {
-          if (stock.isLoco) acc += stock.count;
+          if (isTractionUnit(stock.vehicleRef)) acc += 1;
           return acc;
         }, 0) > 2
       );
     },
 
     teamOnlyVehicles() {
-      return this.store.stockList.filter((stock) => stock.restrictions.teamOnly);
+      return this.store.stockList.filter((stock) => stock.vehicleRef.teamOnly);
     },
 
     hasAnyWarnings() {
@@ -330,6 +342,8 @@ export default defineComponent({
   },
 
   methods: {
+    isTractionUnit,
+
     copyToClipboard() {
       navigator.clipboard.writeText(this.stockString);
 
@@ -346,7 +360,8 @@ export default defineComponent({
       const stock = this.store.stockList[stockID];
 
       this.store.chosenStockListIndex =
-        this.store.chosenStockListIndex == stockID && this.store.chosenVehicle?.type == stock.type
+        this.store.chosenStockListIndex == stockID &&
+        this.store.chosenVehicle?.type == stock.vehicleRef.type
           ? -1
           : stockID;
 
@@ -377,20 +392,6 @@ export default defineComponent({
       this.store.chosenStockListIndex = -1;
     },
 
-    addStock(index: number) {
-      if (index == -1) return;
-
-      this.store.stockList[index].count++;
-    },
-
-    subStock(index: number) {
-      if (index == -1) return;
-
-      if (this.store.stockList[index].count < 2) return;
-
-      this.store.stockList[index].count--;
-    },
-
     removeStock(index: number) {
       if (index == -1) return;
 
@@ -424,7 +425,7 @@ export default defineComponent({
 
     shuffleCars() {
       const availableIndexes = this.store.stockList.reduce((acc, stock, i) => {
-        if (!stock.isLoco) acc.push(i);
+        if (!isTractionUnit(stock.vehicleRef)) acc.push(i);
 
         return acc;
       }, [] as number[]);
@@ -446,7 +447,7 @@ export default defineComponent({
     downloadStock() {
       if (this.store.stockList.length == 0) return alert(this.$t('stocklist.alert-empty'));
 
-      const defaultName = `${this.chosenRealComposition ? this.chosenRealComposition.stockId + ' ' : ''}${this.store.stockList[0].type} ${(this.store.totalWeight / 1000).toFixed(1)}t; ${
+      const defaultName = `${this.chosenRealComposition ? this.chosenRealComposition.stockId + ' ' : ''}${this.store.stockList[0].vehicleRef.type} ${(this.store.totalWeight / 1000).toFixed(1)}t; ${
         this.store.totalLength
       }m; vmax ${this.store.maxStockSpeed}`;
 
@@ -645,7 +646,7 @@ li > .stock-info {
     color: $teamColor;
   }
 
-  &[data-sponsor-only] {
+  &[data-sponsor-only='true'] {
     color: $sponsorColor;
   }
 }
diff --git a/src/components/tabs/WikiListTab.vue b/src/components/tabs/WikiListTab.vue
index ec6a9d7..5edaa30 100644
--- a/src/components/tabs/WikiListTab.vue
+++ b/src/components/tabs/WikiListTab.vue
@@ -48,8 +48,6 @@
           v-for="vehicle in computedVehicles"
           :key="vehicle.type"
           :data-preview="vehicle.type === store.chosenVehicle?.type"
-          :data-sponsor-only="vehicle.restrictions.sponsorOnly > 0"
-          :data-team-only="vehicle.restrictions.teamOnly"
           @click="previewVehicle(vehicle)"
           @dblclick="addVehicle(vehicle)"
           @keydown.enter="previewVehicle(vehicle)"
@@ -58,7 +56,16 @@
           <img loading="lazy" width="120" :src="getThumbnailURL(vehicle.type, 'small')" />
 
           <span>
-            <b class="vehicle-name"> {{ vehicle.type.replace(/_/g, ' ') }} </b>
+            <span
+              class="vehicle-name"
+              :class="{
+                'sponsor-only':
+                  vehicle.sponsorOnlyTimestamp && vehicle.sponsorOnlyTimestamp > Date.now(),
+                'team-only': vehicle.teamOnly,
+              }"
+            >
+              <b>{{ vehicle.type.replace(/_/g, ' ') }}</b>
+            </span>
 
             <div class="vehicle-group">
               {{ $t(`wiki.${vehicle.group}`) }} |
@@ -238,14 +245,6 @@ export default defineComponent({
     background-color: #435288;
   }
 
-  &[data-sponsor-only='true'] {
-    box-shadow: 0 0 5px 0 $sponsorColor;
-  }
-
-  &[data-team-only='true'] {
-    box-shadow: 0 0 5px 0 $teamColor;
-  }
-
   & > span {
     display: flex;
     flex-direction: column;
@@ -260,6 +259,22 @@ export default defineComponent({
   text-overflow: ellipsis;
 }
 
+.sponsor-only {
+  color: $sponsorColor;
+
+  &::after {
+    content: '*';
+  }
+}
+
+.team-only {
+  color: $teamColor;
+
+  &::after {
+    content: '*';
+  }
+}
+
 .vehicle-props {
   color: #ccc;
 }
diff --git a/src/components/utils/StockThumbnails.vue b/src/components/utils/StockThumbnails.vue
index f8e9dbe..c24b38d 100644
--- a/src/components/utils/StockThumbnails.vue
+++ b/src/components/utils/StockThumbnails.vue
@@ -5,7 +5,10 @@
       v-for="(stock, stockIndex) in store.stockList"
       :key="stockIndex"
       :data-selected="store.chosenStockListIndex == stockIndex"
-      :data-sponsor="stock.restrictions.sponsorOnly"
+      :data-sponsor-only="
+        stock.vehicleRef.sponsorOnlyTimestamp && stock.vehicleRef.sponsorOnlyTimestamp > Date.now()
+      "
+      :data-team-only="stock.vehicleRef.teamOnly"
       draggable="true"
       @dragstart="onDragStart(stockIndex)"
       @drop="onDrop($event, stockIndex)"
@@ -13,14 +16,14 @@
       @click="onListItemClick(stockIndex)"
     >
       <b>
-        {{ stock.type }}
+        {{ stock.vehicleRef.type }}
       </b>
 
       <img
         draggable="false"
-        :src="`https://rj.td2.info.pl/dist/img/thumbnails/${stock.type}.png`"
-        :alt="stock.type"
-        :title="stock.type"
+        :src="`https://rj.td2.info.pl/dist/img/thumbnails/${stock.vehicleRef.type}.png`"
+        :alt="stock.vehicleRef.type"
+        :title="stock.vehicleRef.type"
         @error="stockImageError($event, stock)"
       />
     </div>
@@ -43,7 +46,7 @@ const onListItemClick = (index: number) => {
 };
 
 const stockImageError = (e: Event, stock: IStock) => {
-  (e.target as HTMLImageElement).src = `images/${stock.group}-unknown.png`;
+  (e.target as HTMLImageElement).src = `images/${stock.vehicleRef.group}-unknown.png`;
 };
 
 watch(
@@ -88,6 +91,8 @@ const allowDrop = (e: DragEvent) => {
 </script>
 
 <style lang="scss" scoped>
+@import '../../styles/global.scss';
+
 .stock-thumbnails {
   display: flex;
   overflow: auto;
@@ -121,8 +126,12 @@ const allowDrop = (e: DragEvent) => {
     margin: 0 1em;
   }
 
-  &[data-sponsor='true'] > b {
-    color: salmon;
+  &[data-sponsor-only='true'] > b {
+    color: $sponsorColor;
+  }
+
+  &[data-team-only='true'] > b {
+    color: $teamColor;
   }
 
   img {
diff --git a/src/mixins/stockMixin.ts b/src/mixins/stockMixin.ts
index fd80b6f..53e850d 100644
--- a/src/mixins/stockMixin.ts
+++ b/src/mixins/stockMixin.ts
@@ -15,21 +15,11 @@ export default defineComponent({
       return `${Math.random().toString(36).slice(5)}`;
     },
 
-    getStockObject(vehicle: IVehicle, cargo?: ICargo | null, count = 1): IStock {
-      const isLoco = isTractionUnit(vehicle);
-
+    getStockObject(vehicle: IVehicle, cargo?: ICargo | null): IStock {
       return {
         id: this.getStockId(),
-        type: vehicle.type,
-        length: vehicle.length,
-        weight: vehicle.weight,
-        maxSpeed: vehicle.maxSpeed,
-        isLoco,
-        cargo: !isLoco && vehicle.loadable && cargo ? cargo : undefined,
-        count,
-        group: isLoco ? vehicle.group : vehicle.group,
-        constructionType: vehicle.constructionType,
-        restrictions: vehicle.restrictions,
+        vehicleRef: vehicle,
+        cargo: !isTractionUnit(vehicle) && vehicle.loadable && cargo ? cargo : undefined,
       };
     },
 
@@ -38,14 +28,19 @@ export default defineComponent({
 
       const stock = this.getStockObject(vehicle, cargo);
 
-      if (stock.isLoco && !this.store.stockList[0]?.isLoco) this.store.stockList.unshift(stock);
+      if (
+        isTractionUnit(stock.vehicleRef) &&
+        this.store.stockList.length > 0 &&
+        !isTractionUnit(this.store.stockList[0].vehicleRef)
+      )
+        this.store.stockList.unshift(stock);
       else this.store.stockList.push(stock);
     },
 
     addLocomotive(loco: ILocomotive) {
       const stockObj = this.getStockObject(loco);
 
-      if (this.store.stockList.length > 0 && !this.store.stockList[0].isLoco)
+      if (this.store.stockList.length > 0 && !isTractionUnit(this.store.stockList[0].vehicleRef))
         this.store.stockList.unshift(stockObj);
       else this.store.stockList.push(stockObj);
     },
@@ -72,9 +67,9 @@ export default defineComponent({
         let vehicle: IVehicle | null = null;
         let vehicleCargo: ICargo | null = null;
 
-        const isLoco = /^(EU|EP|ET|SM|EN|2EN|SN)/.test(type);
+        const isTractionUnit = /^(EU|EP|ET|SM|EN|2EN|SN)/.test(type);
 
-        if (isLoco) {
+        if (isTractionUnit) {
           const [locoType, spawnProps] = type.split(',');
           vehicle = this.store.locoDataList.find((loco) => loco.type == locoType) || null;
 
diff --git a/src/mixins/stockPreviewMixin.ts b/src/mixins/stockPreviewMixin.ts
index 0f24882..0af938a 100644
--- a/src/mixins/stockPreviewMixin.ts
+++ b/src/mixins/stockPreviewMixin.ts
@@ -14,19 +14,18 @@ export default defineComponent({
 
   methods: {
     previewStock(stock: IStock) {
-      if (stock.isLoco) {
-        const chosenLoco = this.store.locoDataList.find((v) => v.type == stock.type) || null;
-        this.store.chosenVehicle = chosenLoco;
-        this.store.chosenLoco = chosenLoco;
-        this.store.chosenCargo = null;
-        this.store.chosenLocoGroup = stock.group as LocoGroupType;
-      } else {
-        const chosenCar = this.store.carDataList.find((v) => v.type == stock.type) || null;
-        this.store.chosenVehicle = chosenCar;
-        this.store.chosenCar = chosenCar;
+      const vehicleRef = stock.vehicleRef;
 
+      this.store.chosenVehicle = vehicleRef;
+
+      if (isTractionUnit(vehicleRef)) {
+        this.store.chosenLoco = vehicleRef;
+        this.store.chosenCargo = null;
+        this.store.chosenLocoGroup = vehicleRef.group as LocoGroupType;
+      } else {
+        this.store.chosenCar = vehicleRef;
         this.store.chosenCargo = stock.cargo || null;
-        this.store.chosenCarGroup = stock.group as WagonGroupType;
+        this.store.chosenCarGroup = vehicleRef.group as WagonGroupType;
       }
     },
 
diff --git a/src/store.ts b/src/store.ts
index 8bc24da..ccbe398 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -13,6 +13,7 @@ import { defineStore } from 'pinia';
 import {
   acceptableWeight,
   carDataList,
+  isTractionUnit,
   isTrainPassenger,
   locoDataList,
   maxStockSpeed,
@@ -96,25 +97,27 @@ export const useStore = defineStore({
 
     stockSupportsColdStart: (state) => {
       if (state.stockList.length == 0) return false;
-      if (!state.stockList[0].isLoco) return false;
+      if (!isTractionUnit(state.stockList[0].vehicleRef)) return false;
 
       const headingLoco = state.stockList[0];
 
       return (
-        state.vehiclesData?.vehicleProps.find((stock) => stock.type == headingLoco.constructionType)
-          ?.coldStart ?? false
+        state.vehiclesData?.vehicleProps.find(
+          (stock) => stock.type == headingLoco.vehicleRef.constructionType
+        )?.coldStart ?? false
       );
     },
 
     stockSupportsDoubleManning: (state) => {
       if (state.stockList.length == 0) return false;
-      if (!state.stockList[0].isLoco) return false;
+      if (!isTractionUnit(state.stockList[0].vehicleRef)) return false;
 
       const headingLoco = state.stockList[0];
 
       return (
-        state.vehiclesData?.vehicleProps.find((stock) => stock.type == headingLoco.constructionType)
-          ?.doubleManned ?? false
+        state.vehiclesData?.vehicleProps.find(
+          (stock) => stock.type == headingLoco.vehicleRef.constructionType
+        )?.doubleManned ?? false
       );
     },
   },
diff --git a/src/styles/global.scss b/src/styles/global.scss
index d0db630..b5c21fa 100644
--- a/src/styles/global.scss
+++ b/src/styles/global.scss
@@ -6,7 +6,7 @@ $textColor: #fff;
 $secondaryColor: #1b1b1b;
 $accentColor: #e4c428;
 
-$sponsorColor: salmon;
+$sponsorColor: gold;
 $teamColor: #ff4848;
 
 @font-face {
diff --git a/src/types.ts b/src/types.ts
index e8bc3a2..16616d6 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -54,11 +54,12 @@ export interface ILocomotive {
   constructionType: string;
   cabinType: string;
   maxSpeed: number;
-  restrictions: Record<RestrictionType, any>;
   weight: number;
   length: number;
   coldStart: boolean;
   doubleManned: boolean;
+  sponsorOnlyTimestamp: number;
+  teamOnly: boolean;
 }
 
 export interface ICarWagon {
@@ -66,25 +67,18 @@ export interface ICarWagon {
   group: WagonGroupType;
   constructionType: string;
   loadable: boolean;
-  restrictions: Record<RestrictionType, any>;
   maxSpeed: number;
   weight: number;
   length: number;
   cargoTypes: ICargo[];
+  sponsorOnlyTimestamp: number;
+  teamOnly: boolean;
 }
 
 export interface IStock {
   id: string;
-  type: string;
-  group: LocoGroupType | WagonGroupType;
-  constructionType: string;
-  length: number;
-  weight: number;
-  maxSpeed: number;
+  vehicleRef: IVehicle;
   cargo?: ICargo;
-  isLoco: boolean;
-  restrictions: Record<RestrictionType, any>;
-  count: number;
 }
 
 export interface IRealComposition {
diff --git a/src/utils/vehicleUtils.ts b/src/utils/vehicleUtils.ts
index 31b00b5..35608fc 100644
--- a/src/utils/vehicleUtils.ts
+++ b/src/utils/vehicleUtils.ts
@@ -39,7 +39,8 @@ export function locoDataList(vehiclesData: IVehiclesData | undefined) {
       constructionType,
       cabinType,
 
-      restrictions: restrictions ?? {},
+      sponsorOnlyTimestamp: restrictions?.sponsorOnly ?? 0,
+      teamOnly: restrictions?.teamOnly ?? false,
 
       maxSpeed: locoProps.speed,
       length: locoProps.length,
@@ -75,7 +76,8 @@ export function carDataList(vehiclesData: IVehiclesData | undefined) {
       loadable: wagonProps.cargoTypes ? wagonProps.cargoTypes.length > 0 : false,
       cargoTypes: wagonProps?.cargoTypes ?? [],
 
-      restrictions: restrictions ?? {},
+      sponsorOnlyTimestamp: restrictions?.sponsorOnly ?? 0,
+      teamOnly: restrictions?.teamOnly ?? false,
 
       maxSpeed: wagonProps.speed,
       weight: wagonProps?.weight || 0,
@@ -88,25 +90,26 @@ export function carDataList(vehiclesData: IVehiclesData | undefined) {
 
 export function totalWeight(stockList: IStock[]) {
   return stockList.reduce(
-    (acc, stock) => acc + (stock.weight + (stock.cargo?.weight ?? 0)) * stock.count,
+    (acc, stock) => acc + (stock.vehicleRef.weight + (stock.cargo?.weight ?? 0)),
     0
   );
 }
 
 export function totalLength(stockList: IStock[]) {
-  return stockList.reduce((acc, stock) => acc + stock.length * stock.count, 0);
+  return stockList.reduce((acc, stock) => acc + stock.vehicleRef.length, 0);
 }
 
 export function maxStockSpeed(stockList: IStock[]) {
   const stockSpeedLimit = stockList.reduce(
-    (acc, stock) => (stock.maxSpeed < acc || acc == 0 ? stock.maxSpeed : acc),
+    (acc, stock) => (stock.vehicleRef.maxSpeed < acc || acc == 0 ? stock.vehicleRef.maxSpeed : acc),
     0
   );
-  const headingLoco = stockList[0]?.isLoco ? stockList[0] : undefined;
+  const headingLoco =
+    stockList[0] && isTractionUnit(stockList[0].vehicleRef) ? stockList[0] : undefined;
 
   if (!headingLoco) return stockSpeedLimit;
 
-  const locoType = headingLoco.type.split('-')[0];
+  const locoType = headingLoco.vehicleRef.type.split('-')[0];
 
   if (/^(EN|2EN|SN)/.test(locoType)) return stockSpeedLimit;
 
@@ -121,9 +124,9 @@ export function maxStockSpeed(stockList: IStock[]) {
 }
 
 export function acceptableWeight(stockList: IStock[]) {
-  if (stockList.length == 0 || !stockList[0].isLoco) return 0;
+  if (stockList.length == 0 || !isTractionUnit(stockList[0].vehicleRef)) return 0;
 
-  const activeLocomotiveType = stockList[0].type.split('-')[0];
+  const activeLocomotiveType = stockList[0].vehicleRef.type.split('-')[0];
 
   const locoMassLimit = calculateMassLimit(
     activeLocomotiveType as MassLimitLocoType,
@@ -135,9 +138,9 @@ export function acceptableWeight(stockList: IStock[]) {
 
 export function isTrainPassenger(stockList: IStock[]) {
   if (stockList.length == 0) return false;
-  if (stockList.every((stock) => stock.isLoco)) return false;
+  if (stockList.every((stock) => isTractionUnit(stock.vehicleRef))) return false;
 
   return stockList
-    .filter((stock) => !stock.isLoco)
-    .every((stock) => stock.group === 'wagon-passenger');
+    .filter((stock) => !isTractionUnit(stock.vehicleRef))
+    .every((stock) => stock.vehicleRef.group === 'wagon-passenger');
 }