From 820d0bc4e037846a88eb0ff4fcaa96d3346f7103 Mon Sep 17 00:00:00 2001
From: AppleHair <eytanzury@gmail.com>
Date: Wed, 3 Jul 2024 19:44:50 +0300
Subject: [PATCH] [ENHANCEMENT] A better way to limit variations for characters

---
 source/funkin/data/song/SongData.hx        |  5 ++++
 source/funkin/play/song/Song.hx            | 30 ++++++++++++++--------
 source/funkin/ui/freeplay/FreeplayState.hx |  4 +--
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/source/funkin/data/song/SongData.hx b/source/funkin/data/song/SongData.hx
index d2302cd158..ae833a5651 100644
--- a/source/funkin/data/song/SongData.hx
+++ b/source/funkin/data/song/SongData.hx
@@ -63,6 +63,10 @@ class SongMetadata implements ICloneable<SongMetadata>
 
   public var timeChanges:Array<SongTimeChange>;
 
+  @:optional
+  @:default(funkin.util.Constants.DEFAULT_CHARACTER)
+  public var campaignCharacter:String;
+
   /**
    * Defaults to `Constants.DEFAULT_VARIATION`. Populated later.
    */
@@ -86,6 +90,7 @@ class SongMetadata implements ICloneable<SongMetadata>
     this.playData.stage = 'mainStage';
     this.playData.noteStyle = Constants.DEFAULT_NOTE_STYLE;
     this.generatedBy = SongRegistry.DEFAULT_GENERATEDBY;
+    this.campaignCharacter = Constants.DEFAULT_CHARACTER;
     // Variation ID.
     this.variation = (variation == null) ? Constants.DEFAULT_VARIATION : variation;
   }
diff --git a/source/funkin/play/song/Song.hx b/source/funkin/play/song/Song.hx
index dde5ee7b8d..f701838724 100644
--- a/source/funkin/play/song/Song.hx
+++ b/source/funkin/play/song/Song.hx
@@ -80,6 +80,8 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
   // key = variation id, value = metadata
   final _metadata:Map<String, SongMetadata>;
   final difficulties:Map<String, SongDifficulty>;
+  // key = character id, value = allowed variation ids
+  final charVariations:Map<String, Array<String>>;
 
   /**
    * The list of variations a song has.
@@ -150,7 +152,13 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
 
     _data = _fetchData(id);
 
-    _metadata = _data == null ? [] : [Constants.DEFAULT_VARIATION => _data];
+    _metadata = new Map<String, SongMetadata>();
+    charVariations = new Map<String, Array<String>>();
+    if (_data != null)
+    {
+      _metadata.set(Constants.DEFAULT_VARIATION, _data);
+      charVariations.set(_data.campaignCharacter, [Constants.DEFAULT_VARIATION]);
+    }
 
     if (_data != null && _data.playData != null)
     {
@@ -160,6 +168,15 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
         if (variMeta != null)
         {
           _metadata.set(variMeta.variation, variMeta);
+          var variChar:Null<Array<String>> = charVariations.get(variMeta.campaignCharacter);
+          if (variChar != null)
+          {
+            variChar.push(variMeta.variation);
+          }
+          else
+          {
+            charVariations.set(variMeta.campaignCharacter, [variMeta.variation]);
+          }
           trace('  Loaded variation: $vari');
         }
         else
@@ -428,16 +445,7 @@ class Song implements IPlayStateScriptedClass implements IRegistryEntry<SongMeta
   public function getVariationsByCharId(?charId:String):Array<String>
   {
     if (charId == null) charId = Constants.DEFAULT_CHARACTER;
-
-    if (variations.contains(charId))
-    {
-      return [charId];
-    }
-    else
-    {
-      // TODO: How to exclude character variations while keeping other custom variations?
-      return variations;
-    }
+    return charVariations.get(charId) ?? [];
   }
 
   /**
diff --git a/source/funkin/ui/freeplay/FreeplayState.hx b/source/funkin/ui/freeplay/FreeplayState.hx
index 9e32684e54..8821f03cfd 100644
--- a/source/funkin/ui/freeplay/FreeplayState.hx
+++ b/source/funkin/ui/freeplay/FreeplayState.hx
@@ -1576,7 +1576,7 @@ class FreeplayState extends MusicBeatSubState
         FlxG.log.warn('WARN: could not find song with id (${grpCapsules.members[curSelected].songData.songId})');
         return;
       }
-      var targetVariation:String = targetSong.getFirstValidVariation(currentDifficulty);
+      var targetVariation:String = targetSong.getFirstValidVariation(currentDifficulty, targetSong.getVariationsByCharId(currentCharacter));
 
       // TODO: This line of code makes me sad, but you can't really fix it without a breaking migration.
       var suffixedDifficulty = (targetVariation != Constants.DEFAULT_VARIATION
@@ -1717,7 +1717,7 @@ class FreeplayState extends MusicBeatSubState
       return;
     }
     var targetDifficultyId:String = currentDifficulty;
-    var targetVariation:String = targetSong.getFirstValidVariation(targetDifficultyId);
+    var targetVariation:String = targetSong.getFirstValidVariation(targetDifficultyId, targetSong.getVariationsByCharId(currentCharacter));
     PlayStatePlaylist.campaignId = cap.songData.levelId;
 
     var targetDifficulty:SongDifficulty = targetSong.getDifficulty(targetDifficultyId, targetVariation);