From 33964dc99bb46255a763329ba7dd5ecb1eeed8bc Mon Sep 17 00:00:00 2001
From: Anton Vakhrushev <anwinged@ya.ru>
Date: Mon, 12 Oct 2020 13:52:02 +0300
Subject: [PATCH] Fix train troop recurrent tasks

---
 src/Handler/Action/TrainTrooperAction.ts | 15 +++++++++++----
 src/Scheduler.ts                         |  6 ++++--
 src/Village/VillageState.ts              | 19 +++++++++++--------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/src/Handler/Action/TrainTrooperAction.ts b/src/Handler/Action/TrainTrooperAction.ts
index 37796f4..6e08ac1 100644
--- a/src/Handler/Action/TrainTrooperAction.ts
+++ b/src/Handler/Action/TrainTrooperAction.ts
@@ -30,13 +30,20 @@ export class TrainTrooperAction extends BaseAction {
         }
 
         if (nextToTrainCount > 0) {
-            this.scheduler.scheduleTask(TrainTroopTask.name, {
-                ...task.args,
-                count: nextToTrainCount,
-            });
+            this.scheduleNextTask(task, nextToTrainCount);
         }
 
         fillTrainCount(troopId, readyToTrainCount);
         clickTrainButton();
     }
+
+    private scheduleNextTask(prevTask: Task, nextToTrainCount: number) {
+        const villageId = prevTask.args.villageId || taskError('Empty village id');
+        const taskName = prevTask.name;
+        const taskCollection = this.villageFactory.getById(villageId).taskCollection();
+        taskCollection.addTaskAsFirst(taskName, {
+            ...prevTask.args,
+            count: nextToTrainCount,
+        });
+    }
 }
diff --git a/src/Scheduler.ts b/src/Scheduler.ts
index 57ceade..5ef8a4b 100644
--- a/src/Scheduler.ts
+++ b/src/Scheduler.ts
@@ -119,8 +119,10 @@ export class Scheduler {
 
     scheduleTask(name: string, args: Args, ts?: number | undefined): void {
         if (isProductionTask(name) && args.villageId) {
-            const controller = this.villageControllerFactory.getById(args.villageId).controller();
-            controller.addTask(name, args);
+            const taskCollection = this.villageControllerFactory
+                .getById(args.villageId)
+                .taskCollection();
+            taskCollection.addTask(name, args);
         } else {
             this.logger.info('Schedule task', name, args, ts);
             this.taskQueue.add(new Task(uniqTaskId(), ts || timestamp(), name, args));
diff --git a/src/Village/VillageState.ts b/src/Village/VillageState.ts
index f1aef58..6e3a7ac 100644
--- a/src/Village/VillageState.ts
+++ b/src/Village/VillageState.ts
@@ -145,8 +145,8 @@ function createProductionQueueState(
     const taskEndingTimestamp = taskQueueState.finishTs;
     const resources = warehouseState.resources;
     const performance = warehouseState.performance;
-    const firstTaskResources = tasks.slice(0, 1).reduce(taskResourceReducer, Resources.zero());
-    const allTaskResources = tasks.reduce(taskResourceReducer, Resources.zero());
+    const firstTaskResources = getTotalTaskCollectionResources(tasks.slice(0, 1));
+    const allTaskResources = getTotalTaskCollectionResources(tasks);
 
     const currentTimestamp = timestamp();
 
@@ -220,10 +220,14 @@ function getTotalTaskResources(task: TaskCore | undefined): Resources {
     return Resources.zero();
 }
 
-function taskResourceReducer(resources: Resources, task: TaskCore) {
-    return task.args.resources
-        ? resources.add(Resources.fromObject(task.args.resources))
-        : resources;
+function getTotalTaskCollectionResources(tasks: ReadonlyArray<TaskState>): Resources {
+    const reducer = (memo: Resources, task: TaskState): Resources => {
+        const count = task.args.count || 1;
+        return task.args.resources
+            ? memo.add(Resources.fromObject(task.args.resources).scale(count))
+            : memo;
+    };
+    return tasks.reduce(reducer, Resources.zero());
 }
 
 function makeTaskState(task: TaskCore, maxResourcesForTask: Resources): TaskState {
@@ -234,7 +238,6 @@ function makeTaskState(task: TaskCore, maxResourcesForTask: Resources): TaskStat
     const isEnoughGranaryCapacity = maxResourcesForTask.allGreaterOrEqual(
         new Resources(0, 0, 0, taskResources.crop)
     );
-
     const canBeBuilt = isEnoughWarehouseCapacity && isEnoughGranaryCapacity;
 
     return {
@@ -253,7 +256,7 @@ function createVillageState(village: Village, storage: VillageStorage): VillageS
     const capacity = storage.getWarehouseCapacity();
     const performance = storage.getResourcesPerformance();
     const warehouse = makeWarehouseState(resources, capacity, performance);
-    const tasks = storage.getTasks().map((t) => makeTaskState(t, warehouse.optimumFullness));
+    const tasks = storage.getTasks().map((task) => makeTaskState(task, warehouse.optimumFullness));
     const queues = createTaskQueueStates(warehouse, tasks, storage);
     const firstReadyTask = getReadyForProductionTask(queues);
     const firstTaskResources = getTotalTaskResources(firstReadyTask);