#!/usr/bin/env python3
# Startdate: 2024-03-08-6 15:28
# File: srb_lib.py
# Purpose: frontend for srb_lib
# History:
# Usage:
#    tl;dr: ./srb.py --profile 1 --unlock-everything --buy-everything "Profile 1.sav"
# Reference:
#    blog posts 2024-03
#    delsum from github
#    bgconf.py
# WORKHERE: still some confusion around the --unlock-everything --lock-everything and the levelset available levels.

import srb_lib, argparse, sys
from srb_lib import ferror, debuglev

HELP_WEAPONS = ""

for w in srb_lib.WEAPONS:
   HELP_WEAPONS += str(w["id"]) + "," + w["name"] + ','
HELP_WEAPONS = HELP_WEAPONS.rstrip(",")

parser = argparse.ArgumentParser(description="Cli tool for manipulating savegame files for srb.exe",
   formatter_class=argparse.RawDescriptionHelpFormatter,
   epilog=f"""LEVELS include {[i['name'] for i in srb_lib.LEVELS]+list(range(0,len(srb_lib.LEVELS)))}
WEAPONS include {HELP_WEAPONS}
NAME_CHARS include "{srb_lib.NAME_CHARS}"
""")
parser.add_argument("-V","--version",action="version",version="%(prog)s " + srb_lib.srb_lib_version)
parser.add_argument("-d","--debug",nargs='?',default=0,type=int,choices=range(0,11),help="Set debug level")
parser.add_argument("--profile",type=int,choices=range(1,4),help="Profile in user menu.")
parser.add_argument("--get-money",action="store_true",help="Print current money for profile.")
parser.add_argument("--set-money",type=int,help="Set money for profile.")
parser.add_argument("--get-weapon",action="store_true",help="Print currently equipped weapon for profile.")
# choices seems to be too strict here for the numbers. We can live with just the
#choices=[i for i in srb_lib.WEAPONS if i != "undefined"]+list(range(0,16))
#parser.add_argument("--set-weapon",choices=[i for i in srb_lib.WEAPONS if i != "undefined"],help="Print currently equipped weapon for profile.")
parser.add_argument("--set-weapon",help="Set currently equipped weapon for profile.")
parser.add_argument("--get-purchased-weapons",action="store_true",help="Print currently purchased weapons for profile.")
parser.add_argument("--add-purchased-weapons",action="append",help="For profile, add these purchased weapons. Can be used multiple times.")
parser.add_argument("--remove-purchased-weapons",action="append",help="For profile, remove (un-buy) these purchased weapons. Can be used multiple times.")
parser.add_argument("--get-purchased-planes",action="store_true",help="Print currently purchased planes for profile.")
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, balloons, letters for this level for profile. Example value to pass: 0,all 15,none")
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("--set-level-letters",action="append",help="Set collected letters for level for profile. Examples: 0,all 15,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.")
parser.add_argument("--get-tutorial-completed",action="store_true",help="Print if profile has completed the tutorial.")
parser.add_argument("--set-tutorial-completed",choices=["True","False"],help="Set tutorial-completed for profile.")
parser.add_argument("--get-health",action="store_true",help="Print stat health for profile.")
parser.add_argument("--set-health",type=int,choices=range(1,5),help="Set stat health for profile.")
parser.add_argument("--get-stunt",action="store_true",help="Print stat stunt for profile.")
parser.add_argument("--set-stunt",type=int,choices=range(1,5),help="Set stat stunt for profile.")
parser.add_argument("--get-gun",action="store_true",help="Print stat gun for profile.")
parser.add_argument("--set-gun",type=int,choices=range(1,6),help="Set stat gun for profile. 5 is octo-gun which normally unlocks at 100% completion of the game.")
parser.add_argument("--lock-everything",action="store_true",help="Set all levels to completed with rank general, all letters, and all balloons.")
parser.add_argument("--unlock-everything",action="store_true",help="Set all levels to not completed. Take all letters and balloons.")
parser.add_argument("--buy-everything",action="store_true",help="Give all purchasable items: planes, weapons, plane stats.")
parser.add_argument("--unbuy-everything",action="store_true",help="Remove all purchasable items: planes, weapons, plane stats.")
parser.add_argument("--checksum",action=argparse.BooleanOptionalAction,default=True,help="Correct checksum. Default is to do this. It happens at the end of everything else.")
parser.add_argument("file",default="Profile 1.sav")
args = parser.parse_args()
debuglevel = 0
if args.debug:
   debuglevel = args.debug

if debuglev(1,debuglevel):
   ferror("debug level", debuglevel)
if debuglev(8,debuglevel):
   ferror(args)

# common parameters
profile_id = args.profile

#print(f"profile_id={profile_id}")

# 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_status or args.set_level_balloons or args.set_levelset_available_levels or args.set_level_letters or args.unlock_everything or args.lock_everything or args.buy_everything or args.unbuy_everything or args.set_level):
   ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
else:
   if args.get_money:
      money = srb_lib.get_money(args.file, profile_id)
      print(f"Profile {profile_id} has {money} money.")
   if args.set_money:
      if args.set_money > 0xFFFF:
         ferror(f"Warning: Do not set money higher than 65535. While it can technically work, there's no need for that in-game anyways. Continuing...")
      else:
         data = srb_lib.set_money(args.file, profile_id, args.set_money)
         srb_lib.write_file(args.file,0,data)
   if args.get_weapon:
      print(f"Profile {profile_id} has weapon {srb_lib.get_weapon(args.file, profile_id)}")
   if args.set_weapon:
      try:
         args.set_weapon = int(args.set_weapon)
      except:
         pass
      data = srb_lib.set_weapon(args.file, profile_id, args.set_weapon)
      if type(data) == int and data == -1:
         # error is printed in the function
         pass
      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_status:
      for l in args.set_level_status:
         level_num = -1
         try:
            level_num, completion = l.split(",")
         except:
            ferror(f"Warning! Skipping un-parseable level completion spec {l}. Please use \"0,general\" format.")
         if level_num != -1:
            completion_names = [i["name"] for i in srb_lib.LEVEL_STATUSES]
            if completion not in completion_names:
               ferror(f"Warning! Skipping invalid level completion spec {l}; for completion, use one of {completion_names}")
            else:
               # so we have a valid level completion spec.
               data, new_status, message = srb_lib.set_level_status(args.file,profile_id,level_num, completion)
               if (type(data) == int and data == -1) or message != "":
                  ferror(f"Failed to set profile {profile_id} level {level_num} status to {completion} because {message}")
               else:
                  srb_lib.write_file(args.file,0,data)
   if args.set_level_balloons:
      for l in args.set_level_balloons:
         level_num = -1
         try:
            level_num, balloons = l.split(",")
         except:
            ferror(f"Warning! Skipping un-parseable level completion spec {l}. Please use \"0,general\" format.")
         if level_num != -1:
            if balloons not in ["none","all"]:
               ferror(f"Warning! Skipping invalid balloons for level {l}; options are all or none.")
            else:
               data, message = srb_lib.set_level_balloons(args.file,profile_id,level_num,balloons)
               if (type(data) == int and data == -1) or message != "":
                  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:
      print(f"Profile {profile_id} has name {srb_lib.get_name(args.file,profile_id)}")
   if args.set_name:
      data, message = srb_lib.set_name(args.file, profile_id, args.set_name)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to set profile {profile_id} name to {args.set_name} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.get_profile_in_use:
      print(f"Profile {profile_id} in use is {srb_lib.get_profile_in_use(args.file,profile_id)}")
   if args.get_purchased_weapons:
      print(f"Profile {profile_id} has weapons {srb_lib.get_purchased_weapons(args.file,profile_id)}")
   if args.add_purchased_weapons:
      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"add",args.add_purchased_weapons)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to add purchased weapons {args.add_purchased_weapons} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.remove_purchased_weapons:
      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"remove",args.remove_purchased_weapons)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to add purchased weapons {args.remove_purchased_weapons} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.get_tutorial_completed:
      print(f"Profile {profile_id} completed-tutorial is {srb_lib.get_tutorial_completed(args.file,profile_id)}")
   if args.set_tutorial_completed:
      thisbool = False if args.set_tutorial_completed == "False" else True
      data, newstatus = srb_lib.set_tutorial_completed(args.file,profile_id,thisbool)
      srb_lib.write_file(args.file,0,data)
   if args.get_health:
      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"health")
      # error is printed in the function, so if not any error, run
      if not (stat == -1 or message != ""):
         print(f"Profile {profile_id} has health level {stat}")
   if args.set_health:
      data, message = srb_lib.set_plane_stat(args.file,profile_id,"health",args.set_health)
      if data == -1 or message != "":
         print(f"Failed to set health to {args.set_health} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.get_stunt:
      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"stunt")
      # error is printed in the function, so if not any error, run
      if not (stat == -1 or message != ""):
         print(f"Profile {profile_id} has stunt level {stat}")
   if args.set_stunt:
      data, message = srb_lib.set_plane_stat(args.file,profile_id,"stunt",args.set_stunt)
      if data == -1 or message != "":
         print(f"Failed to set stunt to {args.set_stunt} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.get_gun:
      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"gun")
      # error is printed in the function, so if not any error, run
      if not (stat == -1 or message != ""):
         print(f"Profile {profile_id} has gun level {stat}")
   if args.set_gun:
      data, message = srb_lib.set_plane_stat(args.file,profile_id,"gun",args.set_gun)
      if data == -1 or message != "":
         print(f"Failed to set gun to {args.set_gun} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.get_purchased_planes:
      print(f"Profile {profile_id} has planes {srb_lib.get_purchased_planes(args.file,profile_id)}")
   if args.add_purchased_planes:
      data, message = srb_lib.set_purchased_planes(args.file,profile_id,"add",args.add_purchased_planes)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to add purchased planes {args.add_purchased_planes} because {message}")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.remove_purchased_planes:
      data, message = srb_lib.set_purchased_planes(args.file,profile_id,"remove",args.remove_purchased_planes)
      if (type(data) == int and data == -1) or message != "":
         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 levelset available spec {l}. Please use \"0,all\" 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.set_level_letters:
      for l in args.set_level_letters:
         level_num = -1
         try:
            level_num, count = l.split(",")
         except:
            ferror(f"Warning! Skipping unparseable level letter spec {l}. Please use \"0,all\" format.")
         if level_num != -1:
            data, message = srb_lib.set_level_letters(args.file,profile_id,level_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:
               #print(f"good?")
               srb_lib.write_file(args.file,0,data)
   if args.unlock_everything:
      data = srb_lib._get_data_from_data_object(args.file)
      # counting backwards helps the levelset available levels count.
      for i in range(len(srb_lib.LEVELS)-1,0,-1):
         data, new_status, message = srb_lib.set_level_status(data,profile_id,i,"general")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_status, {message}")
            break
         data, message = srb_lib.set_level_balloons(data,profile_id,i,"all")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_balloons, {message}")
            break
         data, message = srb_lib.set_level_letters(data,profile_id,i,"all")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_letters, {message}")
            break
      for i in range(len(srb_lib.LEVELSETS)-1,0,-1):
         data, message = srb_lib.set_levelset_available_levels(data,profile_id,i,"all")
         if message != "":
            ferror(f"Failed on levelset loop {i}, {message}")
            break
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to unlock everything as indicated above.")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.lock_everything:
      data = srb_lib._get_data_from_data_object(args.file)
      # counting backwards helps the levelset available levels count.
      for i in range(len(srb_lib.LEVELS)-1,0,-1):
         data, new_status, message = srb_lib.set_level_status(data,profile_id,i,"none")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_status, {message}")
            break
         data, message = srb_lib.set_level_balloons(data,profile_id,i,"none")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_balloons, {message}")
            break
         data, message = srb_lib.set_level_letters(data,profile_id,i,"none")
         if message != "":
            ferror(f"Failed on loop {i}, set_level_letters, {message}")
            break
      for i in range(len(srb_lib.LEVELSETS)-1,0,-1):
         data, message = srb_lib.set_levelset_available_levels(data,profile_id,i,"none")
         if message != "":
            ferror(f"Failed on levelset loop {i}, {message}")
            break
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to lock everything as indicated above.")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.buy_everything:
      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"add",["all"])
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to remove purchased weapons ['all'] because {message}")
      else:
         data, message = srb_lib.set_purchased_planes(data,profile_id,"add",["all"])
         if (type(data) == int and data == -1) or message != "":
            ferror(f"Failed to add purchased planes ['all'] because {message}")
         else:
            data, message = srb_lib.set_plane_stat(data,profile_id,"health",4)
            if data == -1 or message != "":
               print(f"Failed to set health to 4 because {message}")
            else:
               data, message = srb_lib.set_plane_stat(data,profile_id,"stunt",4)
               if data == -1 or message != "":
                  print(f"Failed to set stunt to 4 because {message}")
               else:
                  data, message = srb_lib.set_plane_stat(data,profile_id,"gun",5)
                  if data == -1 or message != "":
                     print(f"Failed to set gun to 5 because {message}")
      #   srb_lib.write_file(args.file,0,data)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to buy everything as indicated above.")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.unbuy_everything:
      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"remove",["all"])
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to remove purchased weapons ['all'] because {message}")
      else:
         data, message = srb_lib.set_purchased_planes(data,profile_id,"remove",["all"])
         if (type(data) == int and data == -1) or message != "":
            ferror(f"Failed to remove purchased planes ['all'] because {message}")
         else:
            data, message = srb_lib.set_plane_stat(data,profile_id,"health",1)
            if data == -1 or message != "":
               print(f"Failed to set health to 1 because {message}")
            else:
               data, message = srb_lib.set_plane_stat(data,profile_id,"stunt",1)
               if data == -1 or message != "":
                  print(f"Failed to set stunt to 1 because {message}")
               else:
                  data, message = srb_lib.set_plane_stat(data,profile_id,"gun",1)
                  if data == -1 or message != "":
                     print(f"Failed to set gun to 1 because {message}")
      #   srb_lib.write_file(args.file,0,data)
      if (type(data) == int and data == -1) or message != "":
         ferror(f"Failed to unbuy everything as indicated above.")
      else:
         srb_lib.write_file(args.file,0,data)
   if args.set_level:
      for i in args.set_level:
         level_num = -1
         status = ""
         cont = True
         try:
            level_num, status = i.split(",")
         except:
            ferror(f"Warning! Unable to set level spec {i}. Use \"0,all\" or \"15,none\" syntax.")
            cont = False
         if status not in ["all","none"]:
            ferror(f"Warning! Unable to set level {level_num} to {status}. Use all or none.")
            cont = False
         if level_num != -1:
            data, new_status, message = srb_lib.set_level_status(args.file,profile_id,level_num,"general" if status == "all" else "none")
            if message != "":
               ferror(f"Failed on {i}, set_level_status, {message}")
               cont = False
         if cont:
            data, message = srb_lib.set_level_balloons(data,profile_id,level_num,status)
            if message != "":
               ferror(f"Failed on {i}, set_level_balloons, {message}")
               cont = False
         if cont:
            data, message = srb_lib.set_level_letters(data,profile_id,level_num,status)
            if message != "":
               ferror(f"Failed on {i}, set_level_letters, {message}")
               cont = False
         if (type(data) == int and data == -1) or message != "" or not cont:
            ferror(f"Failed to set profile {profile_id} level {level_num} status to {completion} because {message}")
         else:
            srb_lib.write_file(args.file,0,data)

if args.checksum:
   f = args.file
   srb_lib.correct_file(f,debuglevel)