#include "pch.h"
#include "Utils.h"
#include "tiny-AES-c/aes.h"

constexpr BYTE IV[] = {0x72, 0xE0, 0x67, 0xFB, 0xDD, 0xCB, 0xCF, 0x77, 0xEB, 0xE8, 0xBC, 0x64, 0x3F, 0x63, 0x0D, 0x93};
constexpr char podcastKey[] = "\xde" "\xad" "\xbe" "\xef" "\xde" "\xad" "\xbe" "\xef" "\xde" "\xad" "\xbe" "\xef" "\xde"
"\xad" "\xbe" "\xef";
constexpr BYTE nullKey[] = "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0" "\0";

static const std::string urlRegex = "https?:\\/\\/(?:www\.)?([-a-zA-Z0-9@:%._\\+~#=]{1,256}\\.[a-zA-Z0-9()]{1,6})"
"\\b([-a-zA-Z0-9()@:%_\\+.~#?&\\/\\/=]*)";
static const std::string albumSearchPattern = "\x68\x65\x69\x67\x68\x74\x22\x20\x3A\x20\x36\x34\x30";
static const std::wstring songDirRoot = L"Downloads";

bool Utils::Detour32(char* src, char* dst, const intptr_t len)
{
	if (len < 5) return false;

	DWORD  curProtection;
	VirtualProtect(src, len, PAGE_EXECUTE_READWRITE, &curProtection);

	intptr_t  relativeAddress = (intptr_t)(dst - (intptr_t)src) - 5;

	*src = (char)'\xE9';
	*(intptr_t*)((intptr_t)src + 1) = relativeAddress;

	VirtualProtect(src, len, curProtection, &curProtection);
	return true;
}

char* Utils::TrampHook32(char* src, char* dst, const intptr_t len)
{
	// Make sure the length is greater than 5
	if (len < 5) return 0;

	// Create the gateway (len + 5 for the overwritten bytes + the jmp)
	void* gateway = VirtualAlloc(0, len + 5, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);

	if (gateway == NULL) return 0;

	// Write the stolen bytes into the gateway
	memcpy(gateway, src, len);

	// Get the gateway to destination addy
	intptr_t gatewayRelativeAddr = ((intptr_t)src - (intptr_t)gateway) - 5;

	// Add the jmp opcode to the end of the gateway
	*(char*)((intptr_t)gateway + len) = 0xE9;

	// Add the address to the jmp
	*(intptr_t*)((intptr_t)gateway + len + 1) = gatewayRelativeAddr;

	// Perform the detour
	Detour32(src, dst, len);

	return (char*)gateway;
}

int spotifyVer = -1;
int spotifyVerEnd = -1;
int Utils::GetSpotifyVersion()
{
	if (spotifyVer != -1)
		return spotifyVer;

	LPCWSTR lpszFilePath = L"Spotify.exe";
	DWORD dwDummy;
	DWORD dwFVISize = GetFileVersionInfoSize(lpszFilePath, &dwDummy);
	LPBYTE lpVersionInfo = new BYTE[dwFVISize];
	GetFileVersionInfo(lpszFilePath, 0, dwFVISize, lpVersionInfo);
	UINT uLen;
	VS_FIXEDFILEINFO* lpFfi = {};
	VerQueryValue(lpVersionInfo, _T("\\"), (LPVOID*)&lpFfi, &uLen);
	DWORD dwFileVersionMS = lpFfi->dwFileVersionMS;
	DWORD dwFileVersionLS = lpFfi->dwFileVersionLS;
	delete[] lpVersionInfo;

	DWORD dwLeftMost = HIWORD(dwFileVersionMS);
	DWORD dwSecondLeft = LOWORD(dwFileVersionMS);
	DWORD dwSecondRight = HIWORD(dwFileVersionLS);
	DWORD dwRightMost = LOWORD(dwFileVersionLS);

	spotifyVerEnd = dwRightMost;
	return spotifyVer = dwSecondRight;
}

std::string Utils::HexString(BYTE* data, int len)
{
	std::stringstream ss;
	ss << std::hex;

	for (int i(0); i < len; ++i)
		ss << std::setw(2) << std::setfill('0') << (int)data[i];

	return ss.str();
}

std::wstring Utils::FixPathStr(std::wstring str)
{
	// No forbidden characters
	std::wstring badCharsRegex = std::wstring(L"[<>:\"/\\|?*\x00-\x1F]", 14); // Use this constructor for \x00
	str = std::regex_replace(str, std::wregex(badCharsRegex), L"_");

	// No forbidden words or ending periods or spaces
	if (std::regex_match(str, std::wregex(L"^(CON|PRN|AUX|NUL|COM[0-9]|LPT[0-9])$| +$|\\.+$",
		std::regex_constants::ECMAScript | std::regex_constants::icase)))
		str.append(L"_");

	return str;
}

std::wstring Utils::Utf8ToUtf16(const std::string& str)
{
	std::wstring convertedString;
	int requiredSize = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, 0, 0);
	if (requiredSize > 0)
	{
		std::vector<wchar_t> buffer(requiredSize);
		MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, &buffer[0], requiredSize);
		convertedString.assign(buffer.begin(), buffer.end() - 1);
	}

	return convertedString;
}

enum class FileType
{
	OGG,
	MP3
};

struct SongInfo
{
	FileType fileType{ FileType::OGG };
	std::string title{ std::string() };
	std::string artist{ std::string() };
	std::string album{ std::string() };
	std::string coverUrl{ std::string() };
	std::string releaseType{ std::string() };
	std::string releaseDate{ std::string() };
	std::string isrc{ std::string() };
	unsigned int year{ 0 };
	unsigned int trackNum{ 0 };
	unsigned int totalTracks{ 0 };
	unsigned int discNum{ 0 };
	bool isExplicit{ false };
};

void TagSong(std::wstring songPath, SongInfo* songInfo)
{
	TagLib::String fileName = songPath.c_str();
	TagLib::String tTitle(songInfo->title, TagLib::String::Type::UTF8);
	TagLib::String tArtist(songInfo->artist, TagLib::String::Type::UTF8);
	TagLib::String tAlbum(songInfo->album, TagLib::String::Type::UTF8);
	TagLib::String tReleaseType(songInfo->releaseType, TagLib::String::Type::UTF8);
	TagLib::String tReleaseDate(songInfo->releaseDate, TagLib::String::Type::UTF8);
	TagLib::String tIsrc(songInfo->isrc, TagLib::String::Type::UTF8);
	TagLib::String tTotalTracks(std::to_string(songInfo->totalTracks), TagLib::String::Type::UTF8);
	TagLib::String tDiscNum(std::to_string(songInfo->discNum), TagLib::String::Type::UTF8);
	//TagLib::String tIsExplicit(std::to_string(songInfo->isExplicit), TagLib::String::Type::UTF8);

	// Parse episode URL to separate host and path
	std::string coverUrlHost, coverUrlPath;
	try
	{
		std::regex re(urlRegex);
		std::smatch match;
		if (std::regex_search(songInfo->coverUrl, match, re) && match.size() > 1)
		{
			coverUrlHost = match.str(1);
			coverUrlPath = match.str(2);
		}
		else
		{
			std::cout << "Error: Cover art URL is not valid!" << std::endl;
			return;
		}
	}
	catch (std::regex_error& e)
	{
		// Syntax error in the regular expression
		std::cout << "Error: Invalid regex!" << std::endl;
		return;
	}

	std::string coverArtData = Utils::DownloadSpotifyUrl(coverUrlHost, coverUrlPath, "");
	TagLib::ByteVector tCoverArtData(reinterpret_cast<const char*>(&coverArtData[0]), coverArtData.length());

	switch (songInfo->fileType)
	{
	case FileType::OGG:
	{
		TagLib::Ogg::Vorbis::File audioFile(songPath.c_str());
		TagLib::Ogg::XiphComment* tag = audioFile.tag();

		TagLib::FLAC::Picture* coverArt = new TagLib::FLAC::Picture();
		coverArt->setType((TagLib::FLAC::Picture::Type)0x03); // Front Cover
		coverArt->setMimeType("image/jpeg");
		coverArt->setDescription("Front Cover");
		coverArt->setData(tCoverArtData);

		tag->addPicture(coverArt);
		tag->setTitle(tTitle);
		tag->setArtist(tArtist);
		tag->setAlbum(tAlbum);
		tag->setTrack(songInfo->trackNum);
		tag->setYear(songInfo->year);

		tag->addField("DATE", tReleaseDate);
		tag->addField("DISCNUMBER", tDiscNum);
		tag->addField("ISRC", tIsrc);
		tag->addField("SOURCEMEDIA", "Digital Media");
		tag->addField("RELEASETYPE", tReleaseType);
		tag->addField("TOTALTRACKS", tTotalTracks);

		audioFile.save();
		break;
	}
	case FileType::MP3:
	{
		TagLib::MPEG::File audioFile(songPath.c_str());
		TagLib::ID3v2::Tag* tag = audioFile.ID3v2Tag(true);

		tag->setTitle(tTitle);
		tag->setArtist(tArtist);
		tag->setAlbum(tAlbum);
		tag->setTrack(songInfo->trackNum);
		tag->setYear(songInfo->year);

		TagLib::ID3v2::AttachedPictureFrame* frame = new TagLib::ID3v2::AttachedPictureFrame;
		if (frame->picture().size() < tCoverArtData.size())
		{
			frame->setMimeType("image/jpeg");
			frame->setPicture(tCoverArtData);
			tag->addFrame(frame);
		}

		audioFile.save();
		break;
	}
	}
}

int GetUrlNum(int quality)
{
	switch (quality)
	{
	case 0: // Automatic
		return 1; // Set to high quality
	case 1: // Low
		return 8;
	case 2: // Normal
		return 0;
	case 3: // High
		return 1;
	case 4: // Very high
		return 2;
	default:
		return 1; // Shouldn't happen; set to high quality
	}
}

std::string AttemptDecryption(std::string data, std::string key)
{
	// Decrypt encrypted data with Tiny AES in C
	struct AES_ctx ctx;
	AES_init_ctx_iv(&ctx, reinterpret_cast<const uint8_t*>(&key[0]), IV);
	AES_CTR_xcrypt_buffer(&ctx, reinterpret_cast<uint8_t*>(&data[0]), data.size());

	// Check if decrypted song data has valid header (assume Ogg)
	if (data.length() < 4 || data.compare(0, 4, "OggS") != 0)
		return "Error: Not a valid Ogg file after decryption! Please try again";

	return data;
}

std::string DownloadAudioData(std::string fileId, std::string uri, std::string key, std::string authToken, int quality)
{
	std::string audioData;
	std::string songHost, songPath;
	bool isSpotifyHosted;

	// Determine if we need to decrypt the download by parsing fileId for a URL
	try
	{
		std::regex re(urlRegex);
		std::smatch match;
		if (std::regex_search(fileId, match, re) && match.size() > 1)
		{
			// fileId is valid URL (not hosted by Spotify)
			isSpotifyHosted = false;
			songHost = match.str(1);
			songPath = match.str(2);
		}
		else
		{
			// fileId is not valid URL (hosted by Spotify)
			isSpotifyHosted = true;
		}
	}
	catch (std::regex_error& e)
	{
		// Syntax error in the regular expression
		return "Error: Invalid regex while parsing fileID!";
	}

	if (isSpotifyHosted)
	{
		// Get storage resolve from Spotify
		std::string urlNum = std::to_string(GetUrlNum(quality));
		std::string srStr = Utils::DownloadSpotifyUrl("spclient.wg.spotify.com",
			"/storage-resolve/files/audio/interactive_prefetch/" + fileId + "?product=0", authToken);
			//"/storage-resolve/v2/files/audio/interactive/" + urlNum + "/" + fileId + "?product=0", authToken);

		if (srStr.length() <= 5)
			return "Error: Couldn't fetch storage resolve: (" + srStr + ")";
		else if (srStr.substr(0, 5).compare("Error") == 0)
			return srStr;

		// Parse storage resolve response to get the encrypted song data's URL
		std::string songHost;
		std::string songPath;
		try
		{
			std::regex re(urlRegex);
			std::smatch match;
			if (std::regex_search(srStr, match, re) && match.size() > 1)
			{
				songHost = match.str(1);
				songPath = match.str(2);
			}
			else
			{
				return "Error: Download URL not found";
			}
		}
		catch (std::regex_error& e)
		{
			// Syntax error in the regular expression
			return "Error: Invalid regex!";
		}

		// Download encrypted audio data from Spotify
		audioData = Utils::DownloadSpotifyUrl(songHost, songPath, "");

		// Decrypt encrypted audio data using key
		std::string decAudioData = AttemptDecryption(audioData, key);

		if (decAudioData.compare(0, 5, "Error") == 0)
		{
			// Only try again on podcasts
			if (uri.length() > 15 && uri.compare(0, 15, "spotify:episode") == 0)
			{
				std::cout << "Podcast key: " << Utils::HexString((BYTE*)&podcastKey, 16) << std::endl;
				std::string podcastKeyStr = std::string(podcastKey, 16);

				std::cout << "podcastKeyStr = " << podcastKeyStr << std::endl;

				// Try decryption again with podcast key
				decAudioData = AttemptDecryption(audioData, podcastKeyStr);

				if (decAudioData.compare(0, 5, "Error") == 0)
					return "Error: Could not properly decrypt podcast data (try downloading again)!";
			}
			else
				return "Error: Could not properly decrypt song data (try downloading again)!";
		}

		// Remove custom Spotify Ogg page from beginning of file and return audio data
		audioData = decAudioData.substr(audioData.find("\xFF\xFF\xFF\xFFOggS") + 4);
	}
	else
	{
		// Parse episode URL (fileId) to separate host and path
		try
		{
			std::regex re(urlRegex);
			std::smatch match;
			if (std::regex_search(fileId, match, re) && match.size() > 1)
			{
				songHost = match.str(1);
				songPath = match.str(2);
			}
			else
			{
				return "Error: Download URL is not valid!";
			}
		}
		catch (std::regex_error& e)
		{
			// Syntax error in the regular expression
			return "Error: Invalid regex!";
		}

		// Download episode data from URL
		audioData = Utils::DownloadSpotifyUrl(songHost, songPath, "");
	}

	return audioData;
}

static std::wstring songDir = songDirRoot;
void Utils::DownloadSong(std::string fileId, std::string uri, std::string key, std::string authToken, int quality)
{
	std::string downloadStr;
	std::wstring songExtension = L".ogg";

	SongInfo* songInfo = new SongInfo();

	if (fileId.empty() || uri.length() < 13 || authToken.empty())
	{
		std::cout << "Could not download song or episode: missing fileId, trackUri, or authToken!" << std::endl;
		delete songInfo;
		return;
	}

	std::cout << "fileId = " + fileId << std::endl;
	std::cout << "uri =    " + uri << std::endl;
	std::cout << "key =    " + key << std::endl;
	std::cout << "keyHex = " + Utils::HexString((BYTE*)&key[0], 16) << std::endl;

	if (uri.compare(0, 13, "spotify:track") == 0)
	{
		// Not an episode which means no predictable key
		if (memcmp(&key[0], podcastKey, 16) == 0 || memcmp(&key[0], nullKey, 16) == 0)
		{
			std::cout << "Error: Key is of the wrong type! Please try again." << std::endl;
			delete songInfo;
			return;
		}

		std::cout << "Downloading song..." << std::endl;

		// Download song metadata from Spotify API
		std::string metadata = DownloadSpotifyUrl("api.spotify.com", "/v1/tracks/"
			+ uri.substr(uri.find("spotify:track:") + 14), authToken);

		size_t isLocalOff = metadata.find("is_local\" :");

		songInfo->title = strtok((char*)(metadata.substr(metadata.find("name\" :", isLocalOff) + 9)).c_str(), "\"");
		songInfo->artist = strtok((char*)(metadata.substr(metadata.find("name\" :") + 9)).c_str(), "\"");
		songInfo->album = strtok((char*)(metadata.substr(metadata.find(albumSearchPattern) + 404)).c_str(), "\"");
		songInfo->coverUrl = strtok((char*)(metadata.substr(metadata.find("height\" :") + 30)).c_str(), "\"");
		songInfo->releaseType = strtok((char*)(metadata.substr(metadata.find("album_type\" :") + 15)).c_str(), "\"");
		songInfo->releaseDate = strtok((char*)(metadata.substr(metadata.find("release_date\" :") + 17)).c_str(), "\"");
		songInfo->isrc = strtok((char*)(metadata.substr(metadata.find("isrc\" :") + 9)).c_str(), "\"");
		songInfo->trackNum = std::stoi(strtok((char*)(metadata.substr(metadata.find("track_number\" :") + 16)).c_str(),
			","));
		songInfo->totalTracks = std::stoi(strtok((char*)(metadata.substr(metadata.find("total_tracks\" :")
			+ 16)).c_str(), ","));
		songInfo->discNum = std::stoi(strtok((char*)(metadata.substr(metadata.find("disc_number\" :") + 15)).c_str(),
			","));
		songInfo->isExplicit = strtok((char*)(metadata.substr(metadata.find("explicit\" :") + 12)).c_str(), ",");
		songInfo->fileType = FileType::OGG;

		songExtension = L".ogg";
	}
	else if (uri.length() > 15 && uri.compare(0, 15, "spotify:episode") == 0)
	{
		std::string songHost, songPath;

		std::cout << "Downloading episode..." << std::endl;

		// Download episode and show metadata from Spotify API
		std::string metadata = DownloadSpotifyUrl("api.spotify.com", "/v1/episodes/"
			+ uri.substr(uri.find("spotify:episode:") + 16), authToken);

		songInfo->title = strtok((char*)(metadata.substr(metadata.find("name\" :") + 9)).c_str(), "\"");
		songInfo->artist = strtok((char*)(metadata.substr(metadata.find("publisher\" :") + 14)).c_str(), "\"");
		songInfo->album = strtok((char*)(metadata.substr(metadata.find("media_type\" :") + 37)).c_str(), "\"");
		songInfo->coverUrl = strtok((char*)(metadata.substr(metadata.find("height\" :") + 28)).c_str(), "\"");
	}
	else
	{
		std::cout << "Error: Invalid URI!" << std::endl;
		delete songInfo;
		return;
	}

	downloadStr = DownloadAudioData(fileId, uri, key, authToken, quality);

	if (downloadStr.length() < 6)
	{
		std::cout << "Error: Could not download audio!" << std::endl;
		delete songInfo;
		return;
	}
	else if (downloadStr.compare(0, 6, "<HTML>") == 0)
	{
		std::cout << "Error: " + downloadStr << std::endl;
		delete songInfo;
		return;
	}
	else if (downloadStr.compare(0, 5, "Error") == 0)
	{
		std::cout << downloadStr << std::endl;
		delete songInfo;
		return;
	}
	else if (downloadStr.compare(0, 3, "ID3") == 0)
	{
		songInfo->fileType = FileType::MP3;
		songExtension = L".mp3";
	}

	if (songInfo->title.empty() || songInfo->artist.empty() || songInfo->album.empty())
	{
		std::cout << "Error: Empty title/artist/album name!" << std::endl;
		delete songInfo;
		return;
	}

	// TODO: Separate filesystem logic into another function
	std::wstring tempDirArtist = FixPathStr(Utf8ToUtf16(songInfo->artist));

	songDir = songDirRoot;
	if (!CreateDirectoryW(songDir.c_str(), NULL) && ERROR_ALREADY_EXISTS != GetLastError())
		std::cout << "Couldn't create main downloads directory!" << std::endl;
	else if (CreateDirectoryW(std::wstring(songDir + L"\\" + tempDirArtist).c_str(), NULL)
		|| ERROR_ALREADY_EXISTS == GetLastError())
	{
		std::wstring tempDirAlbum = FixPathStr(Utf8ToUtf16(songInfo->album));

		if (CreateDirectoryW(std::wstring(songDir + L"\\" + tempDirArtist + std::wstring(L"\\")
			+ tempDirAlbum).c_str(), NULL) || ERROR_ALREADY_EXISTS == GetLastError())
		{
			std::wstring tempDirSong = FixPathStr(Utf8ToUtf16(songInfo->title));

			songDir += L"\\" + tempDirArtist + std::wstring(L"\\") + tempDirAlbum + L".\\";

			if (songInfo->trackNum != 0)
			{
				std::wstring trackNumStr = std::to_wstring(songInfo->trackNum);
				trackNumStr = std::wstring(2 - trackNumStr.length(), '0') + trackNumStr; // Pad with zeroes

				songDir += trackNumStr + L". ";
			}

			songDir += tempDirArtist + L" - " + tempDirSong + songExtension;

			std::ofstream songFileOut(songDir, std::ios_base::binary);
			songFileOut.write(downloadStr.c_str(), downloadStr.size());
			songFileOut.close();

			TagSong(songDir, songInfo);

			std::cout << "Finished downloading: " << songInfo->artist << " - \"" << songInfo->title << "\"!"
				<< std::endl;

			delete songInfo;
			return;
		}
		else
			std::cout << "Couldn't create album directory!" << std::endl;
	}
	else
		std::cout << "Couldn't create artist directory!" << std::endl;

	delete songInfo;
}

std::string Utils::DownloadSpotifyUrl(std::string host, std::string path, std::string authToken)
{
	std::string response;
	std::string authHeader = (authToken.empty()) ? "" : "Authorization: Bearer " + authToken;
	std::string userAgent = "Spotify/11" + std::to_string(spotifyVer) + std::string("00")
		+ std::to_string(spotifyVerEnd) + std::string(" Win32/Windows 10 (10.0.19042; x64)");
	HINTERNET hSession, hConnect, hRequest;
	BOOL bRequestSent;
	const int bufferSize = 1024;

	hSession = InternetOpenA(userAgent.c_str(), INTERNET_OPEN_TYPE_DIRECT, NULL, NULL, 0);

	if (hSession == NULL)
		return "Error: Could not initialize request: " + GetLastError();

	hConnect = InternetConnectA(hSession, host.c_str(), 80, NULL, NULL, INTERNET_SERVICE_HTTP, 0,
		NULL);

	if (hConnect == NULL)
		return "Error: Could not create connect: " + GetLastError();

	hRequest = HttpOpenRequestA(hConnect, "GET", path.c_str(), NULL, NULL, NULL, INTERNET_FLAG_NO_AUTH, 0);

	if (hRequest == NULL)
		return "Error: Could not create open request: " + GetLastError();

	HttpAddRequestHeadersA(hRequest, authHeader.c_str(), -1, HTTP_ADDREQ_FLAG_ADD | HTTP_ADDREQ_FLAG_REPLACE);
	bRequestSent = HttpSendRequestA(hRequest, NULL, 0, NULL, 0);

	if (!bRequestSent)
		return "Error: Could not send request: " + GetLastError();

	char tmpBuffer[bufferSize] = {};
	BOOL canRead = true;
	DWORD bytesRead = -1;

	while (InternetReadFile(hRequest, tmpBuffer, bufferSize, &bytesRead) && bytesRead)
		response.append(tmpBuffer, bytesRead);

	InternetCloseHandle(hRequest);

	return response;
}

bool Utils::BadPtr(void* ptr)
{
	MEMORY_BASIC_INFORMATION mbi = { 0 };
	if (VirtualQuery(ptr, &mbi, sizeof(mbi)))
	{
		DWORD mask = (PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE
			| PAGE_EXECUTE_WRITECOPY);
		bool b = !(mbi.Protect & mask);

		if (mbi.Protect & (PAGE_GUARD | PAGE_NOACCESS))
			b = true;

		return b;
	}

	return true;
}