#!/usr/bin/env python3
# Startdate: 2024-03-08-6 15:28
# File: srb_lib.py
# Purpose: library for srb.exe savegame hacking
# SPDX-License-Identifier: GPL-3.0-only
# History:
# Usage:
#    from srb.py
# Reference:
#    blog posts 2024-03
#    <https://github.com/8051Enthusiast/delsum>
#    crc width=32 poly=0x4c11db7 init=0x0 xorout=0x235b4b9c refin=false refout=false out_endian=little
#    Incorrect functions: https://gist.github.com/djsweet/3477115595efab31905be7000bb013bc FAILED
#    Fuctions that worked: https://gist.github.com/Lauszus/6c787a3bc26fea6e842dfb8296ebd630 Author: Lauszus
#    https://stackoverflow.com/questions/46109815/reversing-the-byte-order-of-a-string-containing-hexadecimal-characters
#    https://gamefaqs.gamespot.com/pc/930591-snoopy-vs-the-red-baron/faqs/46161
# Documentation:
#    winetricks vd=1024x768
#    winetricks vd=off
# Dependencies:
#    python3-crcmod

import sys, struct
srb_lib_version = "20240309a"

# Table of byte positions of values in the savegame file, minus the first four bytes which are the checksum. Due to zero-indexing of python lists, but for ease of usage, we will always put a zero as the value of index 0. That is, profile 1 will use index 1 of the list.
# money is 0x270 bytes after the "Z<dddddddd" start.
PROFILE_START_POSITION = [0, 0x10, 0x142C, 0x2848]
POS_MONEY = 0x270
POS_EQUIPPED_WEAPON = 0x284

# stinger purchased is relative 0x28C, value 0x40 (b01000000) or 64
# water-balloon-gun is relative 0x2D, value 0x0001
# stinger and potato gun is relative 0x28C, value 0x60, so potato gun is (b00100000) or 64
#POS_STINGER = 0x028C

WEAPONS = [
   "undefined", # 0x0
   "undefined", # secondary weapon is a single machine gun?!
   "undefined",
   "undefined",
   "undefined",
   "Potato Gun", # 0x5
   "Stinger",
   "Woodstock Missile",
   "Water Balloon Cannon",
   "Snow Blower",
   "Fire Boomerang", # 0xA
   "Lightning Rod",
   "Bottle Rockets",
   "Roman Candles",
   "10 Gauge Pumpkin",
   "undefined", # 0xF
]

# WORKHERE: position is for where the balloon info is stored for this levelset, if it is stored here and not per-level.
LEVELSETS = [
   {"id":0,"name":"Aerodrome Island","pos":0x0},
   {"id":1,"name":"Woods of Montsec","pos":0x0},
   {"id":2,"name":"Front Lines of Verdon","pos":0x0},
   {"id":3,"name":"Mines of the Matterhorn","pos":0x0},
   {"id":4,"name":"Verdon Gorge","pos":0x0},
   {"id":5,"name":"Flying Fortress","pos":0x0},
]

# relative to player profile
POS_LEVEL_START = 0x2D8
# each level position is POS_LEVEL_START + (4*level["pos_r"])
# the relative position is because each level set gets 7 (or 6 plus 1 blank) slots, but not all levelsets have 6 levels.
LEVELS = [
   {"id":0, "pos_r":0, "setid":0,"name":"Defend Island"},
   {"id":1, "pos_r":1, "setid":0,"name":"Recover Plans"},
   {"id":2, "pos_r":2, "setid":0,"name":"Protect the Trucks"},
   {"id":3, "pos_r":3, "setid":0,"name":"Attack of the U-Boats"},
   {"id":4, "pos_r":4, "setid":0,"name":"Cripple Outpost Island"},
   {"id":5, "pos_r":5, "setid":0,"name":"Sink the Battleship"},
   {"id":6, "pos_r":7, "setid":1,"name":"Rerun's Challenge"},
   {"id":7, "pos_r":8, "setid":1,"name":"Eliminate Tree Village"},
   {"id":8, "pos_r":9, "setid":1,"name":"Tree Chopper"},
   {"id":9, "pos_r":14,"setid":2,"name":"Trench Warfare"},
   {"id":10,"pos_r":15,"setid":2,"name":"Recover Allied Base"},
   {"id":11,"pos_r":16,"setid":2,"name":"Giant Tank"},
   {"id":12,"pos_r":21,"setid":3,"name":"Surprise Attack"},
   {"id":13,"pos_r":22,"setid":3,"name":"Derail the Train"},
   {"id":14,"pos_r":23,"setid":3,"name":"Enter the Mines"},
   {"id":15,"pos_r":24,"setid":3,"name":"Explore the Mines"},
   {"id":16,"pos_r":25,"setid":3,"name":"Destroy Driller Boss"},
   {"id":17,"pos_r":28,"setid":4,"name":"Navigate the Canyon"},
   {"id":18,"pos_r":29,"setid":4,"name":"Destroy Circus City"},
   {"id":19,"pos_r":30,"setid":4,"name":"Circus Aircraft Carrier"},
   {"id":20,"pos_r":35,"setid":5,"name":"Rescue Allies"},
   {"id":21,"pos_r":36,"setid":5,"name":"Battle the Red Baron"}
]

# hex values
# WOKRHERE: profile 2 level statuses are off somehow?! they are 4 bytes closer to front than the calculation shows.
LEVEL_STATUSES = {
   "59": "all-balloons,all-letters,sergeant",
   "5B": "all-balloons,all-letters,colonel",
   "5D": "all-balloons,all-letters,colonel",
   "60": "all-balloons,all-letters,colonel",
   "62": "all-balloons,all-letters,colonel",
   "63": "all-balloons,all-letters,colonel",
   "64": "all-balloons,all-letters,general"
}

# DEFINE FUNCTIONS
def ferror(*args, **kwargs):
   print(*args, file=sys.stderr, **kwargs)

def debuglev(_numbertocheck, debuglevel):
   # if _numbertocheck <= debuglevel then return truthy
   _debuglev = False
   try:
      if int(_numbertocheck) <= int(debuglevel):
         _debuglev = True
   except Exception as e:
      pass
   return _debuglev

def read_file(filename):
   """ Return the stored checksum and contents of the file """
   with open(filename, 'rb') as f:
      # Read the entire file
      file_contents = f.read()
      checksum = file_contents[0:4]
      file_contents_without_checksum = file_contents[4:]
   return checksum, file_contents_without_checksum

def write_file(filename, checksum, data):
   """ Write the checksum and data to the file. We are assuming the checksum is already correct. """
   print(f"debug: checksum(type={type(checksum).__name__})={checksum}")
   if type(checksum) == int:
      checksum = checksum
   elif type(checksum) == str:
      checksum = int(checksum)
   elif type(checksum) == bytes:
      checksum = int.from_bytes(checksum)
   with open(filename, 'wb') as f:
      f.write(struct.pack('>L',checksum))
      #f.write(struct.pack('>I',int(checksum)))
      f.write(data)

######## start copy-paste 2
#  https://gist.github.com/Lauszus/6c787a3bc26fea6e842dfb8296ebd630

def reflect_data(x, width):
   # See: https://stackoverflow.com/a/20918545
   if width == 8:
      x = ((x & 0x55) << 1) | ((x & 0xAA) >> 1)
      x = ((x & 0x33) << 2) | ((x & 0xCC) >> 2)
      x = ((x & 0x0F) << 4) | ((x & 0xF0) >> 4)
   elif width == 16:
      x = ((x & 0x5555) << 1) | ((x & 0xAAAA) >> 1)
      x = ((x & 0x3333) << 2) | ((x & 0xCCCC) >> 2)
      x = ((x & 0x0F0F) << 4) | ((x & 0xF0F0) >> 4)
      x = ((x & 0x00FF) << 8) | ((x & 0xFF00) >> 8)
   elif width == 32:
      x = ((x & 0x55555555) << 1) | ((x & 0xAAAAAAAA) >> 1)
      x = ((x & 0x33333333) << 2) | ((x & 0xCCCCCCCC) >> 2)
      x = ((x & 0x0F0F0F0F) << 4) | ((x & 0xF0F0F0F0) >> 4)
      x = ((x & 0x00FF00FF) << 8) | ((x & 0xFF00FF00) >> 8)
      x = ((x & 0x0000FFFF) << 16) | ((x & 0xFFFF0000) >> 16)
   else:
      raise ValueError('Unsupported width')
   return x

def crc_poly(data, n, poly, crc=0, ref_in=False, ref_out=False, xor_out=0):
   g = 1 << n | poly  # Generator polynomial
   # Loop over the data
   for d in data:
      # Reverse the input byte if the flag is true
      if ref_in:
         d = reflect_data(d, 8)
      # XOR the top byte in the CRC with the input byte
      crc ^= d << (n - 8)
      # Loop over all the bits in the byte
      for _ in range(8):
         # Start by shifting the CRC, so we can check for the top bit
         crc <<= 1
         # XOR the CRC if the top bit is 1
         if crc & (1 << n):
            crc ^= g
   # Reverse the output if the flag is true
   if ref_out:
      crc = reflect_data(crc, n)
   # Return the CRC value
   return crc ^ xor_out
##### end copy-paste 2

def _get_data_from_data_object(data_object):
   """ Helper function to either open file, or pass bytes through. """
   if type(data_object) == str:
      _, data = read_file(data_object)
   elif type(data_object) == bytes:
      data = data_object
   return data

def _convert_dec_to_reverse_bytes(num_dec):
   """ Helper function to turn a given decimal value into the little-endian values for insertion into the data. """
   # Here is how to do this manually.
   #num_str = str(hex(num_dec))[2:]
   #num_str = ("0" if len(num_str) % 2 else "") + num_str
   #num_bytes = bytearray.fromhex(num_str)
   #num_bytes.reverse()
   #return bytes(num_bytes)
   # I learned how to do this the pythonic way.
   return struct.pack('<i',num_dec)

def get_money(data_object,profile_id):
   """ Print in decimal the current money value of the profile_id. """
   data = _get_data_from_data_object(data_object)
   money_dec = data[PROFILE_START_POSITION[profile_id]+POS_MONEY] + (data[PROFILE_START_POSITION[profile_id]+POS_MONEY+1] * 0x100)
   #money_dec = int(money_hex, base=16)
   return money_dec

def set_money(data_object,profile_id, money_dec):
   """ Using data, set the money given in decimal for the given profile_id. """
   data = _get_data_from_data_object(data_object)
   #money_bytes = _convert_dec_to_reverse_bytes(money_dec)
   #pos_money_low = PROFILE_START_POSITION[profile_id]+POS_MONEY
   data_bytearray = bytearray(data)
   struct.pack_into('<i',data_bytearray,PROFILE_START_POSITION[profile_id]+POS_MONEY,money_dec)
   data = bytes(data_bytearray)
   print(f"after setting money to {money_dec}, we checked and got {get_money(data,profile_id)}")
   return data

def get_weapon(data_object,profile_id):
   data = _get_data_from_data_object(data_object)
   weapon = data[PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON]
   return WEAPONS[weapon]

def set_weapon(data_object,profile_id,weapon):
   data = _get_data_from_data_object(data_object)
   data_bytearray = bytearray(data)
   weapon, weapon_name = get_weapon_info(weapon)
   if weapon == -1:
      ferror(f"Warning: cannot set weapon to {weapon}, because {weapon_name}")
      return -1
   # armed with weapon as the index number, let's change the savegame data
   struct.pack_into('<i',data_bytearray,PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON,weapon)
   data = bytes(data_bytearray)
   print(f"after setting weapon to {WEAPONS[weapon]}, we checked and got {get_weapon(data,profile_id)}")
   return data

def get_weapon_info(weapon):
   """ Returns index, and name of weapon, or -1 and "invalid" """
   if type(weapon) == str:
      try:
         weapon_id = [index for index in range(len(WEAPONS)) if WEAPONS[index].lower() == weapon.lower()][0]
         return weapon_id, WEAPONS[weapon_id]
      except:
         return -1, f"cannot find weapon {weapon}" # must be an incorrect name
   elif type(weapon) == int:
      if weapon in range(0,16):
         return weapon, WEAPONS[weapon]
      else:
         return -1, f"invalid index {weapon}; use 0-15" # must be <0 or >15
   return -1, f"invalid way to reference weapon: [{type(weapon)}]. Use index or name."

def get_level_status(data_object,profile_id,level):
   """ WORKHERE: not sure how to use this, but making sure my position in the data works. """
   data = _get_data_from_data_object(data_object)
   level_obj, message = get_level_info(level)
   if message != "valid" or level == -1:
      ferror(f"Unable to get level status for {level}.")
   print(f"Debug: got level_obj {level_obj} and message {message}")
   profile_level_status = data[PROFILE_START_POSITION[profile_id]+POS_LEVEL_START+(4*level_obj["pos_r"])]
   # it comes back as an int, but does it look better as a hex?
   return hex(profile_level_status)
   
def get_level_info(level):
   """ Returns dictionary of level from LEVELS, searching  by id or name. """
   try:
      # if it is an integer, make sure it shows up as one in the next check
      level = int(level)
   except:
      pass
   if type(level) == str:
      try:
         level = [i for i in LEVELS if i["name"].lower() == level.lower()][0]
         return level, "valid"
      except:
         return -1, f"cannot find level {level}"
   elif type(level) == int:
      if level in range(0,len(LEVELS)+1):
         try:
            level = [i for i in LEVELS if i["id"] == level][0]
            return level, "valid"
         except:
            return -1, f"cannot find level by id {level}"
      else:
         return -1, f"invalid level index {level}; use 0-{len(LEVELS)}"
   return -1, f"invalid way to reference level: [{type(level)}]. Use id or name."

def calculate_checksum(data):
   """ Return the 4-byte checksum used by the game for the provided data. """
   # aka # CRC-32/BZIP2
   # The polynomial appears to be a standard one like used for Ethernet, but I am uncertain of the xor is standard or not.
   # I learned all this only with <https://github.com/8051Enthusiast/delsum>
   # crc width=32 poly=0x4c11db7 init=0x0 xorout=0x235b4b9c refin=false refout=false out_endian=little
   csum = hex(crc_poly(data, 32, 0x4C11DB7, crc=0x0, ref_in=False, ref_out=False, xor_out=0x235b4b9c))
   # comes in as ceb434d4, but needs to be d434b4ce, and bytearray is convenient way to swap it like that
   # trim off the '0x' from the string
   #print(f"Got csum={csum}")
   csum = bytearray.fromhex(str(csum)[2:])
   csum.reverse()
   # and back to bytes because that is how we will want to use it.
   return bytes(csum)

def correct_file(filename, debuglevel = 0):
   """ Given a savegame file, calculate the correct checksum and store that sum instead of whatever is there. """
   cs, data = read_file(filename)
   csum = calculate_checksum(data)
   ran = False
   if csum != cs:
      if debuglev(5,debuglevel):
         ferror(f"Stored checksum is {cs} and will update it to {csum}")
      write_file(filename,csum,data)
      ran = True
   else:
      if debuglev(2,debuglevel):
         ferror(f"Stored checksum is still correct, {csum}. Skipping {filename}")
   return ran