From d1e31500612b9915f69f9b204438ead14f177ff6 Mon Sep 17 00:00:00 2001
From: Luke Hubmayer-Werner <mongoose41@gmail.com>
Date: Sun, 27 Dec 2020 17:51:40 +1030
Subject: [PATCH] Improve chart difficulty selection resilience

---
 scripts/FileLoader.gd  |  6 +++---
 scripts/Library.gd     | 24 ++++++++++++++++++++----
 scripts/Menu.gd        |  2 +-
 scripts/NoteHandler.gd |  4 ++--
 4 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/scripts/FileLoader.gd b/scripts/FileLoader.gd
index 7409212..3ea31d7 100644
--- a/scripts/FileLoader.gd
+++ b/scripts/FileLoader.gd
@@ -590,7 +590,7 @@ func load_folder(folder, filename='song'):
 
 func load_filelist(filelist: Array, directory=''):
 	var charts = {}
-	var key := 1
+	var key := 0
 	for name in filelist:
 		var extension: String = name.rsplit('.', true, 1)[-1]
 		name = directory.rstrip('/') + '/' + name
@@ -600,8 +600,8 @@ func load_filelist(filelist: Array, directory=''):
 				'rgtm':  # multiple charts
 					var res = RGT.load_file(filename)
 					for k in res:
-						charts[k] = res[k]
-				'rgts', 'rgtx':  # single chart
+						charts[Library.difficulty_translations.get(k, k)] = res[k]
+				'rgts', 'rgtx':  # single chart - The keys for this should be translated afterwards
 					charts[key] = RGT.load_file(filename)
 					key += 1
 				'srt':  # maimai, single chart
diff --git a/scripts/Library.gd b/scripts/Library.gd
index 53b6347..6fa9dd9 100644
--- a/scripts/Library.gd
+++ b/scripts/Library.gd
@@ -1,5 +1,7 @@
 extends Node
 
+const difficulty_translations = {'01': 'Z', '02': 'B', '03': 'A', '04': 'E', '05': 'M', '06': 'R', '10': '宴'}  # A bit redundant now but might be useful later for other hacks
+
 class MultilangStr:
 	# Automatically propogate higher langs to lower ones if lower ones are missing.
 	# e.g. if we don't have a proper english title, return the transliterated one instead
@@ -48,7 +50,7 @@ class Song:
 	var video_offsets: Array
 	var audio_preview_times: Array
 	var video_dimensions: Array
-	var chart_difficulties: Dictionary
+	var chart_difficulties := {}
 	const default_difficulty_keys = ['Z', 'B', 'A', 'E', 'M', 'R', '宴']
 	const difficulty_key_ids = {'Z':0, 'B':1, 'A':2, 'E':3, 'M':4, 'R':5, '宴':6}
 
@@ -83,12 +85,10 @@ class Song:
 			TYPE_DICTIONARY:
 				chart_difficulties = diffs
 			TYPE_ARRAY:
-				chart_difficulties = {}
 				for i in min(len(diffs), len(default_difficulty_keys)):
 					chart_difficulties[default_difficulty_keys[i]] = diffs[i]
 			_:
 				print_debug('Invalid chart_difficulties!', title.en)
-				chart_difficulties = {}
 
 	func get_BPM(realtime:=0.0):
 		if not dynamic_bpm:
@@ -128,7 +128,23 @@ func get_song_charts(song_key):
 	if song_key in charts_cache:
 		return charts_cache[song_key]
 	elif song_key in all_songs:
-		charts_cache[song_key] = FileLoader.load_filelist(all_songs[song_key].chart_filelist, 'songs/'+all_songs[song_key].filepath)
+		var charts = FileLoader.load_filelist(all_songs[song_key].chart_filelist, 'songs/'+all_songs[song_key].filepath)
+		# Need to fix keys on this to match the song
+		var diffs = all_songs[song_key].chart_difficulties
+		var missing_diffs = []
+		var unid_charts = []
+		for k in diffs:
+			if not (k in charts):
+				missing_diffs.push_back(k)
+		for k in charts:
+			if not (k in diffs):
+				unid_charts.push_back(k)
+		for i in len(missing_diffs):
+			if i < len(unid_charts):
+				charts[missing_diffs[i]] = charts[unid_charts[i]]
+				charts.erase(unid_charts[i])
+
+		charts_cache[song_key] = charts
 		return charts_cache[song_key]
 	else:
 		print_debug('Invalid song_key: ', song_key)
diff --git a/scripts/Menu.gd b/scripts/Menu.gd
index 96b3677..5dfb93a 100644
--- a/scripts/Menu.gd
+++ b/scripts/Menu.gd
@@ -92,7 +92,7 @@ func _process(delta):
 	if (menu_mode == MenuMode.GAMEPLAY) and (menu_mode_prev_fade_timer <= 0.0) and not NoteHandler.running:
 		var songslist = genres[genres.keys()[selected_genre]]
 		var song_key = songslist[selected_song % len(songslist)]
-		NoteHandler.load_track(song_key, selected_difficulty)
+		NoteHandler.load_track(song_key, Library.Song.default_difficulty_keys[selected_difficulty])
 		NoteHandler.running = true
 
 func draw_string_centered(font, position, string, color := Color.white):
diff --git a/scripts/NoteHandler.gd b/scripts/NoteHandler.gd
index 4d50483..9bd9da5 100644
--- a/scripts/NoteHandler.gd
+++ b/scripts/NoteHandler.gd
@@ -512,14 +512,14 @@ func _ready():
 	InputHandler.connect('button_released', self, 'button_released')
 	InputHandler.connect('touchbutton_released', self, 'touchbutton_released')
 
-func load_track(song_key: String, difficulty_idx: int):
+func load_track(song_key: String, difficulty_key: String):
 	self.song_key = song_key
 	set_time(-3.0)
 	active_notes = []
 	next_note_to_load = 0
 	all_notes = []
 	var data = Library.all_songs[song_key]
-	var chart = Library.get_song_charts(song_key).values()[difficulty_idx]
+	var chart = Library.get_song_charts(song_key)[difficulty_key]
 	for note in chart[1]:
 		all_notes.append(Note.copy_note(note))
 	bpm = data.BPM