WorkplaceScriptableComponent.cs

// Timberborn Mod: Automation
// Author: igor.zavoychinskiy@gmail.com
// License: Public Domain

using System;
using IgorZ.Automation.AutomationSystem;
using IgorZ.Automation.ScriptingEngine.Core;
using IgorZ.Automation.ScriptingEngine.Expressions;
using Timberborn.BaseComponentSystem;
using Timberborn.PrioritySystem;
using Timberborn.WorkSystem;

namespace IgorZ.Automation.ScriptingEngine.ScriptableComponents.Components;

sealed class WorkplaceScriptableComponent : ScriptableComponentBase {

  const string RemoveWorkersActionLocKey = "IgorZ.Automation.Scriptable.Workplace.Action.RemoveWorkers";
  const string SetWorkersActionLocKey = "IgorZ.Automation.Scriptable.Workplace.Action.SetWorkers";
  const string SetPriorityActionLocKey = "IgorZ.Automation.Scriptable.Workplace.Action.SetPriority";
  const string AssignedWorkersSignalLocKey = "IgorZ.Automation.Scriptable.Workplace.Signal.AssignedWorkers";

  const string RemoveWorkersActionName = "Workplace.RemoveWorkers";
  const string SetWorkersActionName = "Workplace.SetWorkers";
  const string SetPriorityActionName = "Workplace.SetPriority";
  const string AssignedWorkersSignalName = "Workplace.AssignedWorkers";

  #region ScriptableComponentBase implementation

  /// <inheritdoc/>
  public override string Name => "Workplace";

  /// <inheritdoc/>
  public override string[] GetSignalNamesForBuilding(AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior, throwIfNotFound: false);
    return workplace ? [AssignedWorkersSignalName] : [];
  }

  /// <inheritdoc/>
  public override Func<ScriptValue> GetSignalSource(string name, AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior);
    return name switch {
        AssignedWorkersSignalName => () => ScriptValue.FromInt(workplace.NumberOfAssignedWorkers),
        _ => throw new UnknownSignalException(name),
    };
  }

  /// <inheritdoc/>
  public override SignalDef GetSignalDefinition(string name, AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior);
    return name switch {
        AssignedWorkersSignalName => LookupSignalDef(
            AssignedWorkersSignalName + "-" + workplace.MaxWorkers,
            () => MakeAssignedWorkersSignalDef(workplace)),
        _ => throw new UnknownSignalException(name),
    };
  }

  /// <inheritdoc/>
  public override void RegisterSignalChangeCallback(SignalOperator signalOperator, ISignalListener host) {
    if (signalOperator.SignalName is not AssignedWorkersSignalName) {
      throw new InvalidOperationException("Unknown signal: " + signalOperator.SignalName);
    }
    host.Behavior.GetOrCreate<WorkplaceChangeTracker>().AddSignal(signalOperator, host);
  }

  /// <inheritdoc/>
  public override void UnregisterSignalChangeCallback(SignalOperator signalOperator, ISignalListener host) {
    host.Behavior.GetOrThrow<WorkplaceChangeTracker>().RemoveSignal(signalOperator, host);
  }

  /// <inheritdoc/>
  public override string[] GetActionNamesForBuilding(AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior, throwIfNotFound: false);
    if (!workplace) {
      return [];
    }
    var workplacePriority = behavior.GetComponentFast<WorkplacePriority>();
    return workplacePriority
        ? [RemoveWorkersActionName, SetWorkersActionName, SetPriorityActionName]
        : [RemoveWorkersActionName, SetWorkersActionName];
  }

  /// <inheritdoc/>
  public override Action<ScriptValue[]> GetActionExecutor(string name, AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior);
    return name switch {
        RemoveWorkersActionName => _ => ResetWorkersAction(workplace),
        SetWorkersActionName => args => SetWorkersAction(workplace, args),
        SetPriorityActionName => args => SetPriorityAction(behavior, args),
        _ => throw new UnknownActionException(name),
    };
  }

  /// <inheritdoc/>
  public override ActionDef GetActionDefinition(string name, AutomationBehavior behavior) {
    var workplace = GetWorkplace(behavior);
    var key = name + "-" + workplace.MaxWorkers;
    return name switch {
        RemoveWorkersActionName => RemoveWorkersActionDef,
        SetWorkersActionName => LookupActionDef(key, () => MakeSetWorkersActionDef(workplace)),
        SetPriorityActionName => SetPriorityActionDef,
        _ => throw new UnknownActionException(name),
    };
  }

  #endregion

  #region Signals

  SignalDef MakeAssignedWorkersSignalDef(Workplace workplace) {
    return new SignalDef {
        ScriptName = AssignedWorkersSignalName,
        DisplayName = Loc.T(AssignedWorkersSignalLocKey),
        Result = new ValueDef {
            ValueType = ScriptValue.TypeEnum.Number,
            ValueFormatter = x => x.AsFloat.ToString("0"),
            ValueValidator = ValueDef.RangeCheckValidatorInt(0, workplace.MaxWorkers),
            ValueUiHint = GetArgumentMinMaxValueHint(0, workplace.MaxWorkers),
        },
    };
  }

  #endregion

  #region Actions

  ActionDef RemoveWorkersActionDef => _removeWorkersActionDef ??= new ActionDef {
      ScriptName = RemoveWorkersActionName,
      DisplayName = Loc.T(RemoveWorkersActionLocKey),
      Arguments = [],
  };
  ActionDef _removeWorkersActionDef;

  ActionDef MakeSetWorkersActionDef(Workplace workplace) {
    return new ActionDef {
        ScriptName = SetWorkersActionName,
        DisplayName = Loc.T(SetWorkersActionLocKey),
        Arguments = [
            new ValueDef {
                ValueType = ScriptValue.TypeEnum.Number,
                ValueFormatter = x => x.AsFloat.ToString("0"),
                ValueValidator = ValueDef.RangeCheckValidatorInt(1, workplace.MaxWorkers),
                ValueUiHint = GetArgumentMinMaxValueHint(1, workplace.MaxWorkers),
            },
        ],
    };
  }

  ActionDef SetPriorityActionDef => _setPriorityActionDef ??= new ActionDef {
      ScriptName = SetPriorityActionName,
      DisplayName = Loc.T(SetPriorityActionLocKey),
      Arguments = [
          new ValueDef {
              ValueType = ScriptValue.TypeEnum.String,
              Options = [
                  ("VeryLow", Loc.T("Priorities.VeryLow")),
                  ("Low", Loc.T("Priorities.Low")),
                  ("Normal", Loc.T("Priorities.Normal")),
                  ("High", Loc.T("Priorities.High")),
                  ("VeryHigh", Loc.T("Priorities.VeryHigh")),
              ],
          },
      ],
  };
  ActionDef _setPriorityActionDef;

  static void ResetWorkersAction(Workplace building) {
    building.DesiredWorkers = 0;
    building.UnassignWorkerIfOverstaffed();
  }

  static void SetWorkersAction(Workplace building, ScriptValue[] args) {
    AssertActionArgsCount(SetWorkersActionName, args, 1);
    var numWorkers = args[0].AsInt;
    if (numWorkers < 1 || numWorkers > building.MaxWorkers) {
      throw new ScriptError.ValueOutOfRange($"Number of workers out of range: {numWorkers}");
    }
    if (building.DesiredWorkers == numWorkers) {
      return;
    }
    building.DesiredWorkers = numWorkers;
    building.UnassignWorkerIfOverstaffed();
  }

  static void SetPriorityAction(AutomationBehavior behavior, ScriptValue[] args) {
    AssertActionArgsCount(SetPriorityActionName, args, 1);
    var priorityName = args[0].AsString;
    if (!Enum.TryParse<Priority>(priorityName, out var priority)) {
      throw new ScriptError.ValueOutOfRange($"Unknown priority: {priorityName}");
    }
    var workplacePriority = behavior.GetComponentFast<WorkplacePriority>();
    if (!workplacePriority) {
      throw new ScriptError.BadStateError(behavior, "Building doesn't have WorkplacePriority");
    }
    if (workplacePriority.Priority == priority) {
      return;
    }
    workplacePriority.SetPriority(priority);
  }

  #endregion

  #region Implementation

  static Workplace GetWorkplace(BaseComponent building, bool throwIfNotFound = true) {
    var workplace = building.GetComponentFast<Workplace>();
    if (!workplace && throwIfNotFound) {
      throw new ScriptError.BadStateError(building, "Building doesn't have Workplace");
    }
    return workplace;
  }

  #endregion

  #region Workplace change tracker

  sealed class WorkplaceChangeTracker : AbstractStatusTracker {

    void Start() {
      var workplace = GetComponentFast<Workplace>();
      workplace.WorkerAssigned += OnWorkerChanged;
      workplace.WorkerUnassigned += OnWorkerChanged;
    }

    void OnWorkerChanged(object sender, WorkerChangedEventArgs args) {
      ScheduleSignal(AssignedWorkersSignalName, ignoreErrors: true);
    }
  }

  #endregion
}