set levelset available levels

2 files changed, 43 insertions, 8 deletions

diff --git a/srb.py b/srb.py
index 9674a66..5d33efe 100755
--- a/srb.py
+++ b/srb.py
@@ -10,9 +10,9 @@
 #    bgconf.py
 
 # WORKHERE:
-# or --set-rank-level 1,general --set-rank-level 1,none --set-rank-level all,general
-# or --give-balloons-level 10 --take-balloons-level 11
 # or --reset-level 5 (which zeros out all info about completion of that level)
+# --set-level-letters 0,none 0,all
+# --set-level 0,done 5,new # sets balloons 10, status general, letters all
 
 import srb_lib, argparse, sys
 from srb_lib import ferror, debuglev
@@ -46,9 +46,10 @@ parser.add_argument("--get-purchased-planes",action="store_true",help="Print cur
 parser.add_argument("--add-purchased-planes",action="append",help="For profile, add these purchased planes. Can be used multiple times.")
 parser.add_argument("--remove-purchased-planes",action="append",help="For profile, remove (un-buy) these purchased planes. Can be used multiple times.")
 parser.add_argument("--get-level",help="Print status for this level for profile.")
-parser.add_argument("--set-level",action="append",help="Set completion status for this level for profile. Example value to pass: \"0,general\"")
+parser.add_argument("--set-level-status",action="append",help="Set completion status for this level for profile. Example value to pass: \"0,general\"")
 parser.add_argument("--set-level-balloons",action="append",help="Set balloon status for this level for profile. Example value to pass: \"0,all\" or \"0,none\"")
 parser.add_argument("--get-levelset",help="Print status for this levelset for profile.")
+parser.add_argument("--set-levelset-available-levels",action="append",help="Set number of available levels in this levelset for profile.")
 parser.add_argument("--get-name",action="store_true",help="Print name for profile.")
 parser.add_argument("--set-name",help="Set name for profile.")
 parser.add_argument("--get-profile-in-use",action="store_true",help="Print if profile is in use.")
@@ -78,7 +79,7 @@ profile_id = args.profile
 #print(f"profile_id={profile_id}")
 
 # WORKHERE: new actions that need --profile must be added here.
-if not profile_id and (args.get_money or args.set_money or args.get_weapon or args.set_weapon or args.get_level or args.get_name or args.set_name or args.get_profile_in_use or args.get_purchased_weapons or args.get_tutorial_completed or args.add_purchased_weapons or args.remove_purchased_weapons or args.get_health or args.get_stunt or args.get_gun or args.set_health or args.set_stunt or args.set_gun or args.get_levelset or args.get_purchased_planes or args.add_purchased_planes or args.remove_purchased_planes or args.set_level or args.set_level_balloons):
+if not profile_id and (args.get_money or args.set_money or args.get_weapon or args.set_weapon or args.get_level or args.get_name or args.set_name or args.get_profile_in_use or args.get_purchased_weapons or args.get_tutorial_completed or args.add_purchased_weapons or args.remove_purchased_weapons or args.get_health or args.get_stunt or args.get_gun or args.set_health or args.set_stunt or args.set_gun or args.get_levelset or args.get_purchased_planes or args.add_purchased_planes or args.remove_purchased_planes or args.set_level_status or args.set_level_balloons or args.set_levelset_available_levels):
    ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
 else:
    if args.get_money:
@@ -105,8 +106,8 @@ else:
          srb_lib.write_file(args.file,0,data)
    if args.get_level:
       print(f"Profile {profile_id} has level {args.get_level} status {srb_lib.get_level_status(args.file,profile_id,args.get_level)}")
-   if args.set_level:
-      for l in args.set_level:
+   if args.set_level_status:
+      for l in args.set_level_status:
          level_num = -1
          try:
             level_num, completion = l.split(",")
@@ -139,7 +140,6 @@ else:
                   ferror(f"Failed to set profile {profile_id} level {level_num} balloons to {balloons} because {message}")
                else:
                   srb_lib.write_file(args.file,0,data)
-
    if args.get_levelset:
       print(f"Profile {profile_id} has levelset {args.get_levelset} status {srb_lib.get_levelset_status(args.file,profile_id,args.get_levelset)}")
    if args.get_name:
@@ -219,6 +219,19 @@ else:
          ferror(f"Failed to add purchased planes {args.remove_purchased_planes} because {message}")
       else:
          srb_lib.write_file(args.file,0,data)
+   if args.set_levelset_available_levels:
+      for l in args.set_levelset_available_levels:
+         levelset_num = -1
+         try:
+            levelset_num, count = l.split(",")
+         except:
+            ferror(f"Warning! Skipping un-parseable level completion spec {l}. Please use \"0,general\" format.")
+         if levelset_num != -1:
+            data, message = srb_lib.set_levelset_available_levels(args.file,profile_id,levelset_num,count)
+            if (type(data) == int and data == -1) or message != "":
+               ferror(f"Failed to set profile {profile_id} level {levelset_num} available levels to {count} because {message}")
+            else:
+               srb_lib.write_file(args.file,0,data)
 
 if args.checksum:
    f = args.file
diff --git a/srb_lib.py b/srb_lib.py
index b659028..64d323e 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -451,10 +451,32 @@ def get_levelset_status(data_object,profile_id,levelset):
       x += 1
    completed_letters = ''.join(completed_letters)
    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_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)
+   if message != "":
+      return -1, f"For set_levelset_available_levels unable to get levelset for {levelset}."
+   if completed_count == "all":
+      completed_count = 8 # no levelset has more than 6 levels so this is a safe maximum, and it will get checked farther below.
+   if completed_count == "none":
+      completed_count = 0
+   # by this point, it better be an integer
+   try:
+      completed_count = int(completed_count)
+   except:
+      return -1, f"cannot set levelset available levels to {completed_count}"
+   if completed_count < 0:
+      completed_count = 0
+   completed_count = min(levelset_obj["l"], completed_count)
+   pos_levelset_completed_mission_mask = PROFILE_START_POSITION[profile_id]+POS_LEVELSET_COMPLETED_MISSIONS_MASK+(INT_SIZE*levelset_obj["id"])
+   completed_bitmask = pow(2,completed_count)-1
+   #print(f"debug: need to set levelset {levelset} available count to {completed_count}, which stored as a bitmask should be {completed_bitmask:07b}")
+   data = srb_pack('<1i',data,pos_levelset_completed_mission_mask,completed_bitmask)
+   return data, ""
+
 def get_collected_balloons_for_levelset(data_object,profile_id,levelset,silent=False):
    data = _get_data_from_data_object(data_object)
    levelset_obj, message = get_levelset_info(levelset)