1 files changed, 19 insertions, 14 deletions

diff --git a/src/music.cpp b/src/music.cpp
index f7498a9..eda3bb8 100644
--- a/src/music.cpp
+++ b/src/music.cpp
@@ -6,6 +6,7 @@ struct HijackedStream {
 	u32 stringOffset;
 	u32 stringOffsetFast;
 	u8 originalID;
+	int streamID;
 };
 
 struct Hijacker {
@@ -98,16 +99,16 @@ const char* SongNameList [] = {
 Hijacker Hijackers[2] = {
 	{
 		{
-			{/*"athletic_lr.n.32.brstm", "athletic_fast_lr.n.32.brstm",*/ _I(0x4A8F8), _I(0x4A938), 4},
-			{/*"BGM_SIRO.32.brstm", "BGM_SIRO_fast.32.brstm",*/ _I(0x4B2E8), _I(0x4B320), 10}
+			{/*"athletic_lr.n.32.brstm", "athletic_fast_lr.n.32.brstm",*/ _I(0x4A8F8), _I(0x4A938), 4, STRM_BGM_ATHLETIC},
+			{/*"BGM_SIRO.32.brstm", "BGM_SIRO_fast.32.brstm",*/ _I(0x4B2E8), _I(0x4B320), 10, STRM_BGM_SHIRO}
 		},
 		0, 0
 	},
 
 	{
 		{
-			{/*"STRM_BGM_CHIJOU.brstm", "STRM_BGM_CHIJOU_FAST.brstm",*/ _I(0x4A83C), _I(0x4A8B4), 1},
-			{/*"STRM_BGM_CHIKA.brstm", "STRM_BGM_CHIKA_FAST.brstm",*/ _I(0x4A878), _I(0x4A780), 2},
+			{/*"STRM_BGM_CHIJOU.brstm", "STRM_BGM_CHIJOU_FAST.brstm",*/ _I(0x4A83C), _I(0x4A8B4), 1, STRM_BGM_CHIJOU},
+			{/*"STRM_BGM_CHIKA.brstm", "STRM_BGM_CHIKA_FAST.brstm",*/ _I(0x4A878), _I(0x4A780), 2, STRM_BGM_CHIKA},
 		},
 		0, 0
 	}
@@ -119,43 +120,46 @@ inline char *BrsarInfoOffset(u32 offset) {
 }
 
 void FixFilesize(u32 streamNameOffset);
+u8 hijackMusicWithSongName(const char *songName, int themeID, bool hasFast, bool useDrums, int *wantRealStreamID);
 
 extern "C" u8 after_course_getMusicForZone(u8 realThemeID) {
 	if (realThemeID < 100)
 		return realThemeID;
 
-	// drums get to use type 1
-	int hjIndex = (realThemeID >= 200) ? 1 : 0;
+	bool usesDrums = (realThemeID >= 200);
+	return hijackMusicWithSongName(SongNameList[realThemeID-100], realThemeID, true, usesDrums, 0);
+}
 
-	Hijacker *hj = &Hijackers[hjIndex];
+u8 hijackMusicWithSongName(const char *songName, int themeID, bool hasFast, bool useDrums, int *wantRealStreamID) {
+	Hijacker *hj = &Hijackers[useDrums ? 1 : 0];
 
 	// do we already have this theme in this slot?
 	// if so, don't switch streams
 	// if we do, NSMBW will think it's a different song, and restart it ...
 	// but if it's just an area transition where both areas are using the same
 	// song, we don't want that
-	if (hj->currentCustomTheme == realThemeID)
+	if ((themeID >= 0) && hj->currentCustomTheme == themeID)
 		return hj->stream[hj->currentStream].originalID;
 
 	// which one do we use this time...?
 	int toUse = (hj->currentStream + 1) & 1;
 
 	hj->currentStream = toUse;
-	hj->currentCustomTheme = realThemeID;
+	hj->currentCustomTheme = themeID;
 
 	// write the stream's info
 	HijackedStream *stream = &hj->stream[hj->currentStream];
 
-	OSReport("%d", realThemeID);
-	OSReport("%s", SongNameList[realThemeID-100]);
-	sprintf(BrsarInfoOffset(stream->stringOffset), "new/%s.er", SongNameList[realThemeID-100]);
-	sprintf(BrsarInfoOffset(stream->stringOffsetFast), "new/%s_F.er", SongNameList[realThemeID-100]);
+	sprintf(BrsarInfoOffset(stream->stringOffset), "new/%s.er", songName);
+	sprintf(BrsarInfoOffset(stream->stringOffsetFast), hasFast?"new/%s_F.er":"new/%s.er", songName);
 
 	// update filesizes
 	FixFilesize(stream->stringOffset);
 	FixFilesize(stream->stringOffsetFast);
 
 	// done!
+	if (wantRealStreamID)
+		*wantRealStreamID = stream->streamID;
 
 	return stream->originalID;
 }
@@ -179,7 +183,8 @@ void FixFilesize(u32 streamNameOffset) {
 			u32 *lengthPtr = (u32*)(streamName - 0x1C);
 			*lengthPtr = info.length;
 		}
-	}
+	} else
+		OSReport("What, I couldn't find \"%s\" :(\n", nameWithSound);
 }