# author: Bluefissure
# License: MIT License
# Description: Fixes Palworld brokwn save files corrupted by someone existing the guild
# Based on the work of https://github.com/cheahjs/palworld-save-tools/releases/tag/v0.13.0

import argparse
import codecs
import os
import json
from lib.gvas import GvasFile
from lib.noindent import CustomEncoder
from lib.palsav import compress_gvas_to_sav, decompress_sav_to_gvas
from lib.paltypes import PALWORLD_CUSTOM_PROPERTIES, PALWORLD_TYPE_HINTS


def main():
    parser = argparse.ArgumentParser(
        prog="palworld-fix-tools",
        description="Fixes Palworld brokwn save files corrupted by someone existing the guild",
    )
    parser.add_argument("filename")
    parser.add_argument(
        "--analyze",
        action="store_true",
        help="Analyzes the file and prints out the missing characters",
    )
    parser.add_argument(
        "--export",
        action="store_true",
        help="Export the content of Level.sav.json",
    )
    parser.add_argument(
        "--fix-missing",
        action="store_true",
        help="Fix the missing players caused by exiting guild by restoring the missing characters from backup",
    )
    parser.add_argument(
        "--fix-capture",
        action="store_true",
        help="Fix the too many capture logs",
    )
    parser.add_argument(
        "--backup",
        help="The backup file to be read from",
    )
    parser.add_argument(
        "--output",
        "-o",
        help="Output file (default: <filename>_fixed.sav)",
    )

    args = parser.parse_args()

    if not os.path.exists(args.filename):
        print(f"{args.filename} does not exist")
        exit(1)

    if not os.path.isfile(args.filename):
        print(f"{args.filename} is not a file")
        exit(1)

    if args.analyze:
        analyze_gvas(args.filename, args.export)
    
    if args.fix_missing:
        if not args.backup:
            print("Backup file is required for fixing")
            exit(1)
        if not args.output:
            output_path = args.filename.replace(".sav", "_fixed.sav")
        else:
            output_path = args.output
        fix_missing(args.filename, args.backup, output_path)
    
    if args.fix_capture:
        if not args.output:
            output_path = args.filename.replace(".sav", "_fixed.sav")
        else:
            output_path = args.output
        fix_capture(args.filename, output_path)
    


def analyze_gvas(filename, export=False) -> (GvasFile, dict, set[str], set[str]):
    print(f"Analyzing {filename}")
    all_players_in_guild = {}
    exist_players_uid = set()
    with open(filename, "rb") as f:
        data = f.read()
        raw_gvas, _ = decompress_sav_to_gvas(data)
    gvas_file = GvasFile.read(raw_gvas, PALWORLD_TYPE_HINTS, PALWORLD_CUSTOM_PROPERTIES)
    if export:
        for (key, value) in gvas_file.properties["worldSaveData"]["value"].items():
            file_path = filename.replace(".sav", "")
            export_file = f"{file_path}_{key}.json"
            print(f"Exporting {key} to {export_file}")
            with codecs.open(export_file, "w", "utf8") as f:
                json.dump(value, f, indent=4, cls=CustomEncoder)
    for group in gvas_file.properties["worldSaveData"]["value"]["GroupSaveDataMap"]["value"]:
        if group["value"]["GroupType"]["value"]["value"] == "EPalGroupType::Guild":
            group_name = group["value"]["RawData"]["value"]["guild_name"]
            group_players = group["value"]["RawData"]["value"]["players"]
            group_players = group["value"]["RawData"]["value"]["players"]
            group_handle_ids = group["value"]["RawData"]["value"]["individual_character_handle_ids"]
            if group_name == "Unnamed Guild" and len(group_players) <= 1:
                continue
            print(f"Analyzing Guild: {group_name} ({len(group_handle_ids)})")
            for player in group_players:
                player_uid = player["player_uid"]
                player_name = player["player_info"]["player_name"]
                all_players_in_guild[player_uid] = player
                print(f"  {player_name}: {player_uid}")
    all_instances = gvas_file.properties["worldSaveData"]["value"]["CharacterSaveParameterMap"]["value"]
    all_instances_uid = set()
    print("Total instances of players/pals/etc: ", len(all_instances))
    for player in all_instances:
        instance_uid = player["key"]["InstanceId"]["value"]
        player_uid = player["key"]["PlayerUId"]["value"]
        all_instances_uid.add(instance_uid)
        if player_uid  == "00000000-0000-0000-0000-000000000000":
            continue
        exist_players_uid.add(player_uid)
    missing_players_uid = set(all_players_in_guild.keys()) - exist_players_uid
    if missing_players_uid:
        print("Missing players:")
        for uid in missing_players_uid:
            player = all_players_in_guild[uid]
            player_name = player["player_info"]["player_name"]
            print(f"  {player_name}: {uid}")
    else:
        print("No missing players")
    return gvas_file, all_players_in_guild, missing_players_uid, all_instances_uid


def fix_missing(filename, backup, output_path):
    if os.path.exists(output_path):
        print(f"{output_path} already exists, this will overwrite the file")
        if not confirm_prompt("Are you sure you want to continue?"):
            exit(1)
    fixed_players = {}
    broken_gvas, all_players_in_guild, missing_players_uid, __ = analyze_gvas(filename)
    if not missing_players_uid:
        print("No missing players, nothing to fix")
        exit(1)

    print(f"Fixing {filename} to {output_path} from backup {backup}")
    with open(backup, "rb") as f:
        data = f.read()
        raw_gvas, _ = decompress_sav_to_gvas(data)
    backup_gvas = GvasFile.read(raw_gvas, PALWORLD_TYPE_HINTS, PALWORLD_CUSTOM_PROPERTIES)
    backup_players = backup_gvas.properties["worldSaveData"]["value"]["CharacterSaveParameterMap"]["value"]
    for player in backup_players:
        player_uid = player["key"]["PlayerUId"]["value"]
        if player_uid in missing_players_uid:
            broken_gvas.properties["worldSaveData"]["value"]["CharacterSaveParameterMap"]["value"].append(player)
            fixed_players[player_uid] = player
    if fixed_players:
        print("Fixed players:")
        for uid, player in fixed_players.items():
            player_name = all_players_in_guild[uid]["player_info"]["player_name"]
            print(f"  {player_name}: {uid}")
    print("Generating SAV file")
    if (
        "Pal.PalWorldSaveGame" in broken_gvas.header.save_game_class_name
        or "Pal.PalLocalWorldSaveGame" in broken_gvas.header.save_game_class_name
    ):
        save_type = 0x32
    else:
        save_type = 0x31
    sav_file = compress_gvas_to_sav(
        broken_gvas.write(PALWORLD_CUSTOM_PROPERTIES), save_type
    )
    print(f"Writing SAV file to {output_path}")
    with open(output_path, "wb") as f:
        f.write(sav_file)

def fix_capture(filename, output_path):
    if os.path.exists(output_path):
        print(f"{output_path} already exists, this will overwrite the file")
        if not confirm_prompt("Are you sure you want to continue?"):
            exit(1)
    (broken_gvas, __, __, all_instances_uid) = analyze_gvas(filename)
    print(f"all_instances_uid: {len(all_instances_uid)}")
    for (idx, group) in enumerate(broken_gvas.properties["worldSaveData"]["value"]["GroupSaveDataMap"]["value"]):
        if group["value"]["GroupType"]["value"]["value"] == "EPalGroupType::Guild":
            group_name = group["value"]["RawData"]["value"]["guild_name"]
            group_handle_ids = group["value"]["RawData"]["value"]["individual_character_handle_ids"]
            temp_instances = []
            for instance in group["value"]["RawData"]["value"]["individual_character_handle_ids"]:
                instance_uid = instance['instance_id']
                if instance_uid in all_instances_uid:
                    temp_instances.append(instance)
            broken_gvas.properties["worldSaveData"]["value"]["GroupSaveDataMap"]["value"][idx]\
                ["value"]["RawData"]["value"]["individual_character_handle_ids"] = temp_instances
            print(f"Fixed capture logs for Guild {group_name}: {len(group_handle_ids)} -> {len(temp_instances)}")
    print("Generating SAV file")
    if (
        "Pal.PalWorldSaveGame" in broken_gvas.header.save_game_class_name
        or "Pal.PalLocalWorldSaveGame" in broken_gvas.header.save_game_class_name
    ):
        save_type = 0x32
    else:
        save_type = 0x31
    sav_file = compress_gvas_to_sav(
        broken_gvas.write(PALWORLD_CUSTOM_PROPERTIES), save_type
    )
    print(f"Writing SAV file to {output_path}")
    with open(output_path, "wb") as f:
        f.write(sav_file)


def confirm_prompt(question: str) -> bool:
    reply = None
    while reply not in ("y", "n"):
        reply = input(f"{question} (y/n): ").casefold()
    return reply == "y"


if __name__ == "__main__":
    main()