add set-level-letters, un/lock everything, un/buy everything

1 files changed, 107 insertions, 15 deletions

diff --git a/srb_lib.py b/srb_lib.py
index 64d323e..933ff6c 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -128,7 +128,7 @@ LEVELS = [
 ]
 
 # hex values
-# WORKHERE: I think balloons and letters merely count as one of the types of requirements. there can up to 5 secondary objectives in a level (maybe the levels without balloons though?). This is probably a bitmask for secondary objectives (including all-balloons,all-letters),health,time.
+# INCOMPLETE: I think balloons and letters merely count as one of the types of requirements. there can up to 5 secondary objectives in a level (maybe the levels without balloons though?). This is probably a bitmask for secondary objectives (including all-balloons,all-letters),health,time.
 # Definitely 0x64 means the whole level is 100% complete, rank General.
 #LEVEL_STATUSES = {
 #   "4A": "3-balloons,no-letters,under-time,good-health,corporal",
@@ -147,6 +147,7 @@ LEVELS = [
 
 # While there are other bitmasks for most ranks, general is only ever 0x64.
 LEVEL_STATUSES = [
+   {"b":0x00,"name":"none"},
    {"b":0x47,"name":"corporal"},
    {"b":0x59,"name":"sergeant"},
    {"b":0x5B,"name":"colonel"},
@@ -351,14 +352,14 @@ def get_purchased_weapons(data_object,profile_id,silent=False):
    return ','.join(weapons_list), weapons_mask
 
 def get_level_status(data_object,profile_id,level,silent=False):
-   """ WORKHERE: need to write set methods for this. """
+   """ Display all sorts of useful info about a level for the profile. The return value is not as useful as the printed info. """
    data = _get_data_from_data_object(data_object)
    level_obj, message = get_level_info(level)
    if message != "" or level == -1:
       ferror(f"Unable to get level status for {level}.")
    levelset_obj, message = get_levelset_info(level_obj["setid"])
    if message != "" or levelset_obj == -1:
-      ferror(f"Unable to get levelset status for {level_obj['setid']}.")
+      ferror(f"Unable to get levelset info for {level_obj['setid']}.")
    if not silent:
       print(f"Debug: for input {level} found level_obj {level_obj} and message {message}")
    pos_level_status = PROFILE_START_POSITION[profile_id]+POS_LEVEL_START+(INT_SIZE*level_obj["pos_r"])
@@ -374,14 +375,13 @@ def get_level_status(data_object,profile_id,level,silent=False):
       print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_level_balloons:04x} collected balloon count: {profile_level_balloons}/{level_obj['b']}")
       print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_level_which_balloons:04x} which balloons: b{profile_level_which_balloons:010b}")
       print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_level_letters:04x} letter count: {profile_level_letters}/{len(level_obj['l'])}")
-   # WORKHERE: show which balloons are collected already for this level?
    # it comes back as an int, but does it look better as a hex?
    return hex(profile_level_status)
 
 def set_level_status(data_object,profile_id,level,status):
    data = _get_data_from_data_object(data_object)
    level_obj, message = get_level_info(level)
-   if message != "" or level == -1:
+   if message != "" or level_obj == -1:
       ferror(f"Unable to get level status for {level}.")
    current_status = get_level_status(data, profile_id, level, silent=True)
    print(f"Before changing, level {level} has status {current_status}")
@@ -390,9 +390,21 @@ def set_level_status(data_object,profile_id,level,status):
       bits = [i for i in LEVEL_STATUSES if i["name"]==status][0]["b"]
    except:
       return data, "", f"unable-to-set-level-completion-status-{status}"
-   print(f"debug: will try to set bits {bits:02x}")
+   #print(f"debug: will try to set bits {bits:02x}")
    data = srb_pack('<1i',data,PROFILE_START_POSITION[profile_id]+POS_LEVEL_START+(INT_SIZE*level_obj["pos_r"]),bits)
    current_status = get_level_status(data, profile_id, level, silent=True)
+   levelset_available_levels = get_levelset_available_levels(data,profile_id,level_obj["setid"])
+   print(f"debug: levelset {level_obj['setid']} currently has {levelset_available_levels} available levels.")
+   # if setting to any completed status, if the levelset available levels is less than this level, then make it this.
+   if levelset_available_levels < (level_obj["set_pos"] + 1) and status not in ["none"]:
+      data, message = set_levelset_available_levels(data,profile_id,level_obj["setid"],level_obj["set_pos"] + 1)
+      if message != "":
+         return -1, -1, f"Unable to set levelset available levels to minimum of this level set_pos {level_obj['set_pos']}"
+   # decrement the available levels in the levelset if clearing out this level and the available levels is exactly this one.
+   if levelset_available_levels == (level_obj["set_pos"] + 1) and status in ["none"]:
+      data, message = set_levelset_available_levels(data,profile_id,level_obj["setid"],level_obj["set_pos"])
+      if message != "":
+         return -1, -1, f"Unable to decrement levelset available levels."
    return data, current_status, ""
 
 def set_level_balloons(data_object,profile_id,level,count):
@@ -419,13 +431,14 @@ def set_level_balloons(data_object,profile_id,level,count):
    data = srb_pack('<1i',data,pos_level_which_balloons,target_which)
    return data, ""
 
-def get_levelset_status(data_object,profile_id,levelset):
-   """ WORKHERE: need to make set methods. """
+def get_levelset_status(data_object,profile_id,levelset,silent=False):
+   """ Display all sorts of useful info about a levelset for the profile. The return value is not as useful as the printed info. """
    data = _get_data_from_data_object(data_object)
    levelset_obj, message = get_levelset_info(levelset)
    if message != "" or levelset == -1:
       ferror(f"Unable to get levelset status for {levelset_obj['setid']}.")
-   print(f"Debug: got levelset {levelset_obj}")
+   if not silent:
+      print(f"Debug: got levelset {levelset_obj}")
    pos_levelset_completed_letters_count = PROFILE_START_POSITION[profile_id]+POS_LEVELSET_COMPLETED_LETTERS_COUNT+(INT_SIZE*2*levelset_obj["id"])
    pos_levelset_completed_letters_mask = pos_levelset_completed_letters_count + INT_SIZE
    pos_levelset_completed_mission_mask = PROFILE_START_POSITION[profile_id]+POS_LEVELSET_COMPLETED_MISSIONS_MASK+(INT_SIZE*levelset_obj["id"])
@@ -439,21 +452,100 @@ def get_levelset_status(data_object,profile_id,levelset):
    collected_balloons_for_levelset, message = get_collected_balloons_for_levelset(data,profile_id,levelset,silent=True)
    if message != "":
       ferror("Unable to count collected balloons.")
-   print(f"Debug: levelset {levelset_obj['id']},\"{levelset_obj['name']:23s}\" has mission mask: b{pow(2,levelset_obj['l'])-1:06b}")
-   print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_mission_mask:04x} levelset_obj {levelset_obj['id']}          completed missions: {profile_levelset_status}")
-   print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_letters_count:04x} completed letter count: {profile_levelset_completed_letters_count}/{len(levelset_obj['le'])}")
-   print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_letters_mask:04x} completed letter mask: b{profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}")
-   print(f"Debug: collected balloons: {collected_balloons_for_levelset}/{levelset_obj['b']} {levelset_obj['c']}")
+   if not silent:
+      print(f"Debug: levelset {levelset_obj['id']},\"{levelset_obj['name']:23s}\" has mission mask: b{pow(2,levelset_obj['l'])-1:06b}")
+      print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_mission_mask:04x} levelset_obj {levelset_obj['id']}          completed missions: {profile_levelset_status}")
+      print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_letters_count:04x} completed letter count: {profile_levelset_completed_letters_count}/{len(levelset_obj['le'])}")
+      print(f"Debug: abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_letters_mask:04x} completed letter mask: b{profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}")
+      print(f"Debug: collected balloons: {collected_balloons_for_levelset}/{levelset_obj['b']} {levelset_obj['c']}")
    completed_letters = list(levelset_obj['le'])
    x = 0
    for i in f"{profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}":
       completed_letters[x] = completed_letters[x].upper() if int(i) == 1 else completed_letters[x]
       x += 1
    completed_letters = ''.join(completed_letters)
-   print(f"Debug:  letters, uppercase is collected:  {completed_letters}")
+   if not silent:
+      print(f"Debug:  letters, uppercase is collected:  {completed_letters}")
    # it comes back as an int, but does it look better as a hex?
    return profile_levelset_status
 
+def set_level_letters(data_object,profile_id,level,letters):
+   """ Set collected letters for given level to all or none. """
+   data = _get_data_from_data_object(data_object)
+   level_obj, message = get_level_info(level)
+   if message != "" or level == -1:
+      return -1, f"Unable to get level status for {level}."
+   levelset_obj, message = get_levelset_info(level_obj["setid"])
+   if message != "":
+      return -1, f"For set_level_letters unable to get levelset for {level_obj['setid']}."
+   pos_level_letters = PROFILE_START_POSITION[profile_id]+POS_LEVEL_LETTERS+(INT_SIZE*level_obj["set_pos"])+(INT_SIZE*POS_LEVEL_LETTERS_MULTIPLIER_LEVELSET*level_obj["setid"])
+   if letters not in ["all","none"]:
+      return -1, f"Cannot set level letters to {letters}. Use all or none."
+   if letters == "all":
+      letter_count = len(level_obj["l"])
+   else:
+      letter_count = 0
+   # update letter count for level.
+   data = srb_pack('<1i',data,pos_level_letters,letter_count)
+   pos_levelset_completed_letters_count = PROFILE_START_POSITION[profile_id]+POS_LEVELSET_COMPLETED_LETTERS_COUNT+(INT_SIZE*2*levelset_obj["id"])
+   #print(f"debug: setting abs 0x{CHECKSUM_LENGTH+pos_levelset_completed_letters_count:04x} to {letter_count}")
+   #data = srb_pack('<1i',data,pos_levelset_completed_letters_count,letter_count)
+   # prepare to update the levelset completed letters count and mask
+   pos_levelset_completed_letters_mask = pos_levelset_completed_letters_count + INT_SIZE
+   # need to get current letter bitmask for levelset
+   # convert the left-1-bitshifted value to useful, left-to-right mask
+   profile_levelset_completed_letters_mask = struct.unpack_from('<1i',data,pos_levelset_completed_letters_mask)[0] >> 1
+   # python trick to reverse a custom-formatted string and convert back to int while reading the string as base 2
+   profile_levelset_completed_letters_mask = int(f"{profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}"[::-1],2)
+   print(f"debug: levelset {levelset_obj['id']} before changes has letters mask {profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}")
+   levels_to_check = [i for i in LEVELS if i["setid"] == levelset_obj["id"]]
+   #print(f"need to check levels {levels_to_check}")
+   levelset_letters_count = 0
+   x = 0
+   bitmask_pos = 0
+   while x < len(levels_to_check):
+      ltc = [i for i in levels_to_check if i["set_pos"] == x][0]
+      this_level_count = data[PROFILE_START_POSITION[profile_id]+POS_LEVEL_LETTERS+(INT_SIZE*ltc["set_pos"])+(INT_SIZE*POS_LEVEL_LETTERS_MULTIPLIER_LEVELSET*ltc["setid"])]
+      #print(f"debug: level {ltc['id']} has count {this_level_count}")
+      levelset_letters_count += this_level_count
+      if ltc["id"] == level_obj["id"]:
+         this_bitmask_pos = bitmask_pos
+      bitmask_pos += len(ltc["l"])
+      x += 1
+   # so now that this_bitmask_pos is populated with the location in the levelset letter mask of this level's letters, we need to manipulate it.
+   #print(f"debug: need to manipulate letter mask {profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b} at position {this_bitmask_pos} length {len(level_obj['l'])}")
+   new_letter_mask = list(f"{profile_levelset_completed_letters_mask:0{len(levelset_obj['le'])}b}")
+   #print(f"debug: this_bitmask_pos={this_bitmask_pos},len(level_obj['l'])={len(level_obj['l'])}, thing={(['1'] if letters == 'all' else ['0'])*len(level_obj['l'])}")
+   # I could not get this pythonic thing working the way I wanted. I think the slice manipulation that uses len(level_obj["l"]) was not working....
+   #new_letter_mask[this_bitmask_pos:len(level_obj["l"])+1] = (["1"] if letters == "all" else ["0"]) * len(level_obj["l"])
+   # ... so I just run it per spot individually.
+   x = 0
+   while x < len(level_obj["l"]):
+      new_letter_mask[this_bitmask_pos+x] = "1" if letters == "all" else "0"
+      x += 1
+   #print(f"debug: intermediate, new_letter_mask={new_letter_mask}")
+   new_letter_mask = "".join(new_letter_mask)
+   new_letter_mask_int = int(new_letter_mask,2)
+   #print(f"debug: new_letter_mask updated list is str {new_letter_mask}, int {new_letter_mask_int}")
+   #print(f"debug: total letters collected for levelset {levelset_obj['id']}: {levelset_letters_count}")
+   print(f"debug: new letters bitmask for levelset {levelset_obj['id']}: {new_letter_mask_int:0{len(new_letter_mask)}b}")
+   # bitmask needs to be reversed and left-1-bitshifted
+   new_letter_mask_final = int(f"{new_letter_mask_int:0{len(levelset_obj['le'])}b}"[::-1],2) << 1
+   print(f"debug: so that final bitmask should be {new_letter_mask_final:016b}")
+   # set levelset letter count
+   data = srb_pack('<1i',data,pos_levelset_completed_letters_count,levelset_letters_count)
+   # set levelset letter mask
+   data = srb_pack('<1i',data,pos_levelset_completed_letters_mask,new_letter_mask_final)
+   return data, ""
+
+def get_levelset_available_levels(data_object,profile_id,levelset):
+   data = _get_data_from_data_object(data_object)
+   levelset_obj, message = get_levelset_info(levelset)
+   if message != "":
+      return -1, f"For set_levelset_available_levels unable to get levelset for {levelset}."
+   pos_levelset_completed_mission_mask = PROFILE_START_POSITION[profile_id]+POS_LEVELSET_COMPLETED_MISSIONS_MASK+(INT_SIZE*levelset_obj["id"])
+   return data[pos_levelset_completed_mission_mask]
+
 def set_levelset_available_levels(data_object,profile_id,levelset,completed_count):
    data = _get_data_from_data_object(data_object)
    levelset_obj, message = get_levelset_info(levelset)