Add: AI API for vehicle group colours (#7336)

4 files changed, 72 insertions, 0 deletions

diff --git a/src/script/api/ai/ai_group.hpp.sq b/src/script/api/ai/ai_group.hpp.sq
index 5fa50ae1b..29d86754d 100644
--- a/src/script/api/ai/ai_group.hpp.sq
+++ b/src/script/api/ai/ai_group.hpp.sq
@@ -45,6 +45,10 @@ void SQAIGroup_Register(Squirrel *engine)
 	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::GetProfitThisYear,           "GetProfitThisYear",           2, ".i");
 	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::GetProfitLastYear,           "GetProfitLastYear",           2, ".i");
 	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::GetCurrentUsage,             "GetCurrentUsage",             2, ".i");
+	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::SetPrimaryColour,            "SetPrimaryColour",            3, ".ii");
+	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::SetSecondaryColour,          "SetSecondaryColour",          3, ".ii");
+	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::GetPrimaryColour,            "GetPrimaryColour",            2, ".i");
+	SQAIGroup.DefSQStaticMethod(engine, &ScriptGroup::GetSecondaryColour,          "GetSecondaryColour",          2, ".i");
 
 	SQAIGroup.PostRegister(engine);
 }
diff --git a/src/script/api/ai_changelog.hpp b/src/script/api/ai_changelog.hpp
index f3ec539a4..373590801 100644
--- a/src/script/api/ai_changelog.hpp
+++ b/src/script/api/ai_changelog.hpp
@@ -19,6 +19,12 @@
  *
  * This version is not yet released. The following changes are not set in stone yet.
  *
+ * API additions:
+ * \li AIGroup::SetPrimaryColour
+ * \li AIGroup::SetSecondaryColour
+ * \li AIGroup::GetPrimaryColour
+ * \li AIGroup::GetSecondaryColour
+ *
  * \b 1.9.0
  *
  * API additions:
diff --git a/src/script/api/script_group.cpp b/src/script/api/script_group.cpp
index 16a50bc17..fcae6c3d2 100644
--- a/src/script/api/script_group.cpp
+++ b/src/script/api/script_group.cpp
@@ -194,3 +194,35 @@
 
 	return occupancy / vehicle_count;
 }
+
+/* static */ bool ScriptGroup::SetPrimaryColour(GroupID group_id, ScriptCompany::Colours colour)
+{
+	EnforcePrecondition(false, IsValidGroup(group_id));
+
+	return ScriptObject::DoCommand(0, group_id, colour << 16, CMD_SET_GROUP_LIVERY);
+}
+
+/* static */ bool ScriptGroup::SetSecondaryColour(GroupID group_id, ScriptCompany::Colours colour)
+{
+	EnforcePrecondition(false, IsValidGroup(group_id));
+
+	return ScriptObject::DoCommand(0, group_id, (1 << 8) | (colour << 16), CMD_SET_GROUP_LIVERY);
+}
+
+/* static */ ScriptCompany::Colours ScriptGroup::GetPrimaryColour(GroupID group_id)
+{
+	EnforcePrecondition(ScriptCompany::Colours::COLOUR_INVALID, IsValidGroup(group_id));
+
+	const Group *g = ::Group::GetIfValid(group_id);
+	if (!HasBit(g->livery.in_use, 0)) return ScriptCompany::Colours::COLOUR_INVALID;
+	return (ScriptCompany::Colours)g->livery.colour1;
+}
+
+/* static */ ScriptCompany::Colours ScriptGroup::GetSecondaryColour(GroupID group_id)
+{
+	EnforcePrecondition(ScriptCompany::Colours::COLOUR_INVALID, IsValidGroup(group_id));
+
+	const Group *g = ::Group::GetIfValid(group_id);
+	if (!HasBit(g->livery.in_use, 1)) return ScriptCompany::Colours::COLOUR_INVALID;
+	return (ScriptCompany::Colours)g->livery.colour2;
+}
diff --git a/src/script/api/script_group.hpp b/src/script/api/script_group.hpp
index 9e88eaff1..6e7deb0b7 100644
--- a/src/script/api/script_group.hpp
+++ b/src/script/api/script_group.hpp
@@ -214,6 +214,36 @@ public:
 	 * @return The current usage of the group.
 	 */
 	static uint32 GetCurrentUsage(GroupID group_id);
+
+	/**
+	 * Set primary colour for a group.
+	 * @param group_id The group id to set the colour of.
+	 * @param colour Colour to set.
+	 * @pre IsValidGroup(group_id).
+	 */
+	static bool SetPrimaryColour(GroupID group_id, ScriptCompany::Colours colour);
+
+	/**
+	 * Set secondary colour for a group.
+	 * @param group_id The group id to set the colour of.
+	 * @param colour Colour to set.
+	 * @pre IsValidGroup(group_id).
+	 */
+	static bool SetSecondaryColour(GroupID group_id, ScriptCompany::Colours colour);
+
+	/**
+	 * Get primary colour of a group.
+	 * @param group_id The group id to get the colour of.
+	 * @pre IsValidGroup(group_id).
+	 */
+	static ScriptCompany::Colours GetPrimaryColour(GroupID group_id);
+
+	/**
+	 * Get secondary colour for a group.
+	 * @param group_id The group id to get the colour of.
+	 * @pre IsValidGroup(group_id).
+	 */
+	static ScriptCompany::Colours GetSecondaryColour(GroupID group_id);
 };
 
 #endif /* SCRIPT_GROUP_HPP */