import json
import os
import subprocess
import sys
import zlib

UESAVE_TYPE_MAPS = [
    ".worldSaveData.CharacterSaveParameterMap.Key=Struct",
    ".worldSaveData.FoliageGridSaveDataMap.Key=Struct",
    ".worldSaveData.FoliageGridSaveDataMap.ModelMap.InstanceDataMap.Key=Struct",
    ".worldSaveData.MapObjectSpawnerInStageSaveData.Key=Struct",
    ".worldSaveData.ItemContainerSaveData.Key=Struct",
    ".worldSaveData.CharacterContainerSaveData.Key=Struct",
]

def main():
    if len(sys.argv) < 5:
        print('fix-host-save.py <uesave.exe> <save_path> <player_a_guid> <player_b_guid>')
        exit(1)
    
    uesave_path = sys.argv[1]
    save_path = sys.argv[2]
    player_a_guid = sys.argv[3]
    player_b_guid = sys.argv[4]
    host_guid = '00000000000000000000000000000001'
    
    
    # Users accidentally include the .sav file extension when copying the GUID over. Only the GUID should be passed.
    if player_a_guid[-4:] == '.sav' or player_b_guid[-4:] == '.sav':
        print('ERROR: It looks like you\'re providing the whole name of the file instead of just the GUID. For example, instead of using "<GUID>.sav" in the command, you should be using only the GUID.')
        exit(1)
    
    # Users accidentally remove characters from their GUIDs when copying it over. All GUIDs should be 32 characters long.
    if len(player_a_guid) != 32:
        print('ERROR: Your <new_guid> should be 32 characters long, but it is ' + str(len(player_a_guid)) + ' characters long. Make sure you copied the exact GUID.')
        exit(1)
    
    if len(player_b_guid) != 32:
        print('ERROR: Your <player_b_guid> should be 32 characters long, but it is ' + str(len(player_b_guid)) + ' characters long. Make sure you copied the exact GUID.')
        exit(1)
    
    # Apply expected formatting for the GUID.
    player_a_guid_formatted = '{}-{}-{}-{}-{}'.format(player_a_guid[:8], player_a_guid[8:12], player_a_guid[12:16], player_a_guid[16:20], player_a_guid[20:]).lower()
    host_guid_formatted = '{}-{}-{}-{}-{}'.format(host_guid[:8], host_guid[8:12], host_guid[12:16], host_guid[16:20], host_guid[20:]).lower()
    
    level_sav_path = save_path + '/Level.sav'
    player_a_sav_path = save_path + '/Players/' + host_guid + '.sav'
    player_b_sav_path = save_path + '/Players/'+ player_b_guid + '.sav'
    level_json_path = level_sav_path + '.json'
    player_a_json_path = player_a_sav_path + '.json'
    player_b_json_path = player_b_sav_path + '.json'


    # uesave_path must point directly to the executable, not just the path it is located in.
    if not os.path.exists(uesave_path) or not os.path.isfile(uesave_path):
        print('ERROR: Your given <uesave_path> of "' + uesave_path + '" is invalid. It must point directly to the executable. For example: C:\\Users\\Bob\\.cargo\\bin\\uesave.exe')
        exit(1)
    
    # save_path must exist in order to use it.
    if not os.path.exists(save_path):
        print('ERROR: Your given <save_path> of "' + save_path + '" does not exist. Did you enter the correct path to your save folder?')
        exit(1)
    
    # Player B needs to have created a character on the co-op server and that save is used for this script.
    if not os.path.exists(player_b_sav_path):
        print('ERROR: Player B\'s save does not exist. Did you enter the correct GUID of your player? It should look like "8E910AC2000000000000000000000000".\nDid your player create their character with the provided save? Once they create their character, a file called "' + player_b_sav_path + '" should appear. Look back over the steps in the README on how to get your GUID.')
        exit(1)
    
    # Warn the user about potential data loss.
    print('WARNING: Running this script WILL change your save files and could \
potentially corrupt your data. It is HIGHLY recommended that you make a backup \
of your save folder before continuing. Press enter if you would like to continue.')
    input('> ')
    
    # Convert save files to JSON so it is possible to edit them.
    print('Converting save files to JSON...', flush=True)
    sav_to_json(uesave_path, level_sav_path)
    sav_to_json(uesave_path, player_a_sav_path)
    sav_to_json(uesave_path, player_b_sav_path)
    print('Done!', flush=True)
    
    # Parse our JSON files.
    print('Parsing JSON files...', end='', flush=True)
    with open(player_a_json_path) as f:
        player_a_json = json.load(f)
    with open(player_b_json_path) as f:
        player_b_json = json.load(f)
    with open(level_json_path) as f:
        level_json = json.load(f)
    print('Done!', flush=True)
    
    # Replace all instances of the host's GUID with the player A's GUID, and the player B's GUID with the host's GUID.
    print('Modifying JSON save data...', end='', flush=True)
    
    # Player data replacement : host <- A.
    player_a_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = player_a_guid_formatted
    player_a_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["IndividualId"]["Struct"]["value"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = player_a_guid_formatted
    player_a_instance_id = player_a_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["IndividualId"]["Struct"]["value"]["Struct"]["InstanceId"]["Struct"]["value"]["Guid"]

    # Player data replacement : B <- host.
    player_b_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = host_guid_formatted
    player_b_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["IndividualId"]["Struct"]["value"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = host_guid_formatted
    player_b_instance_id = player_b_json["root"]["properties"]["SaveData"]["Struct"]["value"]["Struct"]["IndividualId"]["Struct"]["value"]["Struct"]["InstanceId"]["Struct"]["value"]["Guid"]
    
    # Level data replacement.
    instance_ids_len = len(level_json["root"]["properties"]["worldSaveData"]["Struct"]["value"]["Struct"]["CharacterSaveParameterMap"]["Map"]["value"])
    for i in range(instance_ids_len):
        a_done = b_done = False
        instance_id = level_json["root"]["properties"]["worldSaveData"]["Struct"]["value"]["Struct"]["CharacterSaveParameterMap"]["Map"]["value"][i]["key"]["Struct"]["Struct"]["InstanceId"]["Struct"]["value"]["Guid"]
        if instance_id == player_a_instance_id:
            level_json["root"]["properties"]["worldSaveData"]["Struct"]["value"]["Struct"]["CharacterSaveParameterMap"]["Map"]["value"][i]["key"]["Struct"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = player_a_guid_formatted
            a_done = True
        if instance_id == player_b_instance_id:
            level_json["root"]["properties"]["worldSaveData"]["Struct"]["value"]["Struct"]["CharacterSaveParameterMap"]["Map"]["value"][i]["key"]["Struct"]["Struct"]["PlayerUId"]["Struct"]["value"]["Guid"] = host_guid_formatted
            b_done = True
        if a_done and b_done:
            break
    
    print('Done!', flush=True)
    
    # Dump modified data to JSON.
    print('Exporting JSON data...', end='', flush=True)
    with open(player_a_json_path, 'w') as f:
        json.dump(player_a_json, f, indent=2)
    with open(player_b_json_path, 'w') as f:
        json.dump(player_b_json, f, indent=2)
    with open(level_json_path, 'w') as f:
        json.dump(level_json, f, indent=2)
    print('Done!')
    
    # Convert our JSON files to save files.
    print('Converting JSON files back to save files...', flush=True)
    json_to_sav(uesave_path, level_json_path)
    json_to_sav(uesave_path, player_a_json_path)
    json_to_sav(uesave_path, player_b_json_path)
    print('Done!', flush=True)
    
    # Clean up miscellaneous GVAS and JSON files which are no longer needed.
    print('Cleaning up miscellaneous files...', end='', flush=True)
    clean_up_files(level_sav_path)
    clean_up_files(player_a_sav_path)
    clean_up_files(player_b_sav_path)
    print('Done!', flush=True)
    
    # We must rename the patched save files from the host GUID to player A's GUID and player B's GUID to the host GUID for the server to recognize them.
    os.rename(player_a_sav_path, save_path + '/Players/' + player_a_guid + '.sav')
    os.rename(player_b_sav_path, player_a_sav_path)
    print('Fix has been applied! Have fun!')

def sav_to_json(uesave_path, file):
    with open(file, 'rb') as f:
        # Read the file
        data = f.read()
        uncompressed_len = int.from_bytes(data[0:4], byteorder='little')
        compressed_len = int.from_bytes(data[4:8], byteorder='little')
        magic_bytes = data[8:11]
        save_type = data[11]
        # Check for magic bytes
        if magic_bytes != b'PlZ':
            print(f'File {file} is not a save file, found {magic_bytes} instead of P1Z')
            return
        # Valid save types
        if save_type not in [0x30, 0x31, 0x32]:
            print(f'File {file} has an unknown save type: {save_type}')
            return
        # We only have 0x31 (single zlib) and 0x32 (double zlib) saves
        if save_type not in [0x31, 0x32]:
            print(f'File {file} uses an unhandled compression type: {save_type}')
            return
        if save_type == 0x31:
            # Check if the compressed length is correct
            if compressed_len != len(data) - 12:
                print(f'File {file} has an incorrect compressed length: {compressed_len}')
                return
        # Decompress file
        uncompressed_data = zlib.decompress(data[12:])
        if save_type == 0x32:
            # Check if the compressed length is correct
            if compressed_len != len(uncompressed_data):
                print(f'File {file} has an incorrect compressed length: {compressed_len}')
                return
            # Decompress file
            uncompressed_data = zlib.decompress(uncompressed_data)
        # Check if the uncompressed length is correct
        if uncompressed_len != len(uncompressed_data):
            print(f'File {file} has an incorrect uncompressed length: {uncompressed_len}')
            return
        # Save the uncompressed file
        with open(file + '.gvas', 'wb') as f:
            f.write(uncompressed_data)
        print(f'File {file} uncompressed successfully')
        # Convert to json with uesave
        # Run uesave.exe with the uncompressed file piped as stdin
        # Standard out will be the json string
        uesave_run = subprocess.run(uesave_to_json_params(uesave_path, file+'.json'), input=uncompressed_data, capture_output=True)
        # Check if the command was successful
        if uesave_run.returncode != 0:
            print(f'uesave.exe failed to convert {file} (return {uesave_run.returncode})')
            print(uesave_run.stdout.decode('utf-8'))
            print(uesave_run.stderr.decode('utf-8'))
            return
        print(f'File {file} (type: {save_type}) converted to JSON successfully')

def json_to_sav(uesave_path, file):
    # Convert the file back to binary
    gvas_file = file.replace('.sav.json', '.sav.gvas')
    sav_file = file.replace('.sav.json', '.sav')
    uesave_run = subprocess.run(uesave_from_json_params(uesave_path, file, gvas_file))
    if uesave_run.returncode != 0:
        print(f'uesave.exe failed to convert {file} (return {uesave_run.returncode})')
        return
    # Open the old sav file to get type
    with open(sav_file, 'rb') as f:
        data = f.read()
        save_type = data[11]
    # Open the binary file
    with open(gvas_file, 'rb') as f:
        # Read the file
        data = f.read()
        uncompressed_len = len(data)
        compressed_data = zlib.compress(data)
        compressed_len = len(compressed_data)
        if save_type == 0x32:
            compressed_data = zlib.compress(compressed_data)
        with open(sav_file, 'wb') as f:
            f.write(uncompressed_len.to_bytes(4, byteorder='little'))
            f.write(compressed_len.to_bytes(4, byteorder='little'))
            f.write(b'PlZ')
            f.write(bytes([save_type]))
            f.write(bytes(compressed_data))
    print(f'Converted {file} to {sav_file}')

def clean_up_files(file):
    os.remove(file + '.json')
    os.remove(file + '.gvas')

def uesave_to_json_params(uesave_path, out_path):
    args = [
        uesave_path,
        'to-json',
        '--output', out_path,
    ]
    for map_type in UESAVE_TYPE_MAPS:
        args.append('--type')
        args.append(f'{map_type}')
    return args

def uesave_from_json_params(uesave_path, input_file, output_file):
    args = [
        uesave_path,
        'from-json',
        '--input', input_file,
        '--output', output_file,
    ]
    return args

if __name__ == "__main__":
    main()