#!/usr/bin/env python3

from argparse import ArgumentParser, FileType
from hashlib import md5
from requests import get
from struct import pack
from urllib.parse import urlencode

# char magic[8]
# u32 game_id
# u32 console_id
# u32 forum_topic_id
# u32 flags
# u32 is_final
# u32 achievement_count
# u32 leaderboard_count
# u32 reserved1[3]
# u8 hash[16]
# char title[32]
# char publisher[32]
# char developer[32]
# char genre[32]
# char release_date[32]
# char console_name[32]
HEADER_STRUCT = "<8sIIIIIII12x16s32s32s32s32s32s32s"  # 0x100 bytes

# u32 id
# u32 points
# u32 flags
# u32 creation_time
# u32 modified_time
# u32 description_len
# u32 mem_addr_len
# u32 completion_time
# u32 reserved1[8]
# char title[32]
# char author[32]
# char reserved2[128]
# char description[256]
# char mem_addr[512]
ACHIEVEMENT_STRUCT = "<IIIIIII4x32x32s32s128x256s512s"  # 0x400 bytes

# u32 id
# u32 hidden
# u32 lower_is_better
# u32 description_len
# u32 mem_len
# u32 reserved1[11]
# char title[32]
# char format[32]
# char reserved2[128]
# char description[256]
# char mem[512]
LEADERBOARD_STRUCT = "<IIIII44x32s32s128x256s512s"  # 0x400 bytes

def log(str):
	if args.verbose:
		print(str)

def rc_hash_nintendo_ds(input):
	offset = 0

	input.seek(0)
	header = input.read(512)
	if len(header) != 512:
		raise Exception("Failed to read header")

	if header[0] == 0x2E and header[1] == 0x00 and header[2] == 0x00 and header[3] == 0xEA and \
		header[0xB0] == 0x44 and header[0xB1] == 0x46 and header[0xB2] == 0x96 and header[0xB3] == 0:
		# SuperCard header detected, ignore it
		log("Ignoring SuperCard header")

		offset = 512
		input.seek(offset)
		header = input.read(512)
		if len(header) != 512:
			raise Exception("Failed to read header")

	arm9_addr = header[0x20] | (header[0x21] << 8) | (header[0x22] << 16) | (header[0x23] << 24)
	arm9_size = header[0x2C] | (header[0x2D] << 8) | (header[0x2E] << 16) | (header[0x2F] << 24)
	arm7_addr = header[0x30] | (header[0x31] << 8) | (header[0x32] << 16) | (header[0x33] << 24)
	arm7_size = header[0x3C] | (header[0x3D] << 8) | (header[0x3E] << 16) | (header[0x3F] << 24)
	icon_addr = header[0x68] | (header[0x69] << 8) | (header[0x6A] << 16) | (header[0x6B] << 24)

	if arm9_size + arm7_size > 16 * 1024 * 1024:
		# sanity check - code blocks are typically less than 1MB each - assume not a DS ROM
		raise Exception("arm9 code size (%u) + arm7 code size (%u) exceeds 16MB" % (arm9_size, arm7_size))

	log("Hashing 352 byte header")
	hash_buffer = header[:0x160]

	log("Hashing %u byte arm9 code (at %08X)" % (arm9_size, arm9_addr))
	input.seek(arm9_addr + offset)
	hash_buffer += input.read(arm9_size)

	log("Hashing %u byte arm7 code (at %08X)" % (arm7_size, arm7_addr))
	input.seek(arm7_addr + offset)
	hash_buffer += input.read(arm7_size)

	log("Hashing 2560 byte icon and labels data (at %08X)" % icon_addr)
	input.seek(icon_addr + offset)
	icon_buf = input.read(0xA00)
	if len(icon_buf) < 0xA00:
		# some homebrew games don't provide a full icon block, and no data after the icon block.
		# if we didn't get a full icon block, fill the remaining portion with 0s
		log("Warning: only got %u bytes for icon and labels data, 0-padding to 2560 bytes" % len(icon_buf))
		icon_buf += b"\0" * (0xA00 - len(icon_buf))
	hash_buffer += icon_buf

	with open("test.bin", "wb") as f:
		f.write(hash_buffer)

	return md5(hash_buffer)

def api_get(request, vars={}):
	vars["r"] = request
	j = get("http://retroachievements.org/dorequest.php?" + urlencode(vars)).json()
	if j["Success"]:
		return j
	else:
		raise Exception("API request failed")


def bsach(input, output, username, password):
	hash = rc_hash_nintendo_ds(input)
	log("Hash: %s" % hash.hexdigest())
	game_id = api_get("gameid", {"m": hash})["GameID"]
	log("Game ID: %u" % game_id)
	token = api_get("login", {"u": username, "p": password})["Token"]
	patch_data = api_get("patch", {"u": "Pk11", "t": token, "g": 12711})["PatchData"]

	output.write(pack(HEADER_STRUCT,
		b"BSACHv1",
		patch_data["ID"],
		patch_data["ConsoleID"],
		patch_data["ForumTopicID"],
		patch_data["Flags"],
		patch_data["IsFinal"],
		len(patch_data["Achievements"]),
		len(patch_data["Leaderboards"]),
		hash.digest(),
		patch_data["Title"].encode("utf8"),
		patch_data["Publisher"].encode("utf8"),
		patch_data["Developer"].encode("utf8"),
		patch_data["Genre"].encode("utf8"),
		patch_data["Released"].encode("utf8"),
		patch_data["ConsoleName"].encode("utf8")
	))

	for achievement in patch_data["Achievements"]:
		output.write(pack(ACHIEVEMENT_STRUCT,
			achievement["ID"],
			achievement["Points"],
			achievement["Flags"],
			achievement["Created"],
			achievement["Modified"],
			len(achievement["Description"]),
			len(achievement["MemAddr"]),
			achievement["Title"].encode("utf8"),
			achievement["Author"].encode("utf8"),
			achievement["Description"].encode("utf8"),
			achievement["MemAddr"].encode("utf8")
		))

	for leaderboard in patch_data["Leaderboards"]:
		output.write(pack(LEADERBOARD_STRUCT,
			leaderboard["ID"],
			leaderboard["Hidden"],
			leaderboard["LowerIsBetter"],
			len(leaderboard["Description"]),
			len(leaderboard["Mem"]),
			leaderboard["Title"].encode("utf8"),
			leaderboard["Format"].encode("utf8"),
			leaderboard["Description"].encode("utf8"),
			leaderboard["Mem"].encode("utf8")
		))

	log("%s successfully created with %d achievements and %d leaderboards" % (output.name, len(patch_data["Achievements"]), len(patch_data["Leaderboards"])))

if __name__ == "__main__":
	parser = ArgumentParser(description="Downloads RetroAchievements to an nds-bootstrap friendly format")
	parser.add_argument("input", type=FileType("rb"), help="input nds")
	parser.add_argument("output", type=FileType("wb"), help="output bsach")
	parser.add_argument("username", type=str, help="RetroAchievements username")
	parser.add_argument("password", type=str, help="RetroAchievements password")
	parser.add_argument("-v", "--verbose", action="store_true", help="enable args.verbose logging")

	args = parser.parse_args()

	bsach(args.input, args.output, args.username, args.password)