Rough script to merge Rocksmith profiles based on higest played count in Learn a Song

profilemerge.py

""" This is a quick and dirty tool to merge two Rocksmith Profiles. 
It is very hacky and not terribly safe. Use with care.

Merging is done by choosing the arrangement with the highest "Learn a song" play count.

For a given arrangement and play count: 
    If both files have the same play count, nothing changes.
    If the master file has a higher play count than the merge file, nothing changes.
    If the merge file has a higher play count, the learn a song data, the arrangement 
    stats, and the score attack data are copied to the master file. 

The steps are:
- Select the master file which changes will be merged into. 
- Select the merge data file.
- Do some basic checks on consistency of arrangements between the two files.
- Do the merge. 
- Write the file.
"""
# cspell: ignore Playnexts

import argparse

from pathlib import Path

from typing import List, Any

import simplejson

from rsrtools.files.exceptions import RsrError
from rsrtools.files.profilemanager import RSProfileManager


class ProfileArrData:
    def __init__(self, json_data: Any) -> None:
        self.learn_a_song = json_data["Songs"]
        self.play_next = json_data["Playnexts"]["Songs"]
        self.score_attack = json_data["SongsSA"]
        self.stats = json_data["Stats"]["Songs"]
        # won't import streak_id
        streak_id = json_data["Achievements"]["SongStreakID"]

        # primitive check on set keys and internal consistency
        arr_ids = (
            self.learn_a_song.keys()
            | self.play_next.keys()
            | self.score_attack.keys()
            | self.stats.keys()
        )
        arr_ids.union(streak_id)

        # self checks - extremely inefficient, but should hardly ever be used.
        json_str = simplejson.dumps(json_data).upper()
        for arr in arr_ids:
            if not arr.isupper():
                raise RsrError("Unexpected lower case arrangement id.")

            arr_count = json_str.count(arr.upper())
            check = 0
            for song_dict in (
                self.learn_a_song,
                self.score_attack,
                self.play_next,
                self.stats,
            ):
                if arr in song_dict:
                    check = check + 1

            if arr == streak_id:
                check = check + 1

            if arr_count != check:
                # I don't believe this should happen, but this will catch arrangement
                # locations I'm not aware of.
                raise RsrError(
                    f"Expected arrangement f{arr} in to appear f{arr_count} times, "
                    f"but found f{check} instances."
                )

    def transfer_arr(self, arr_id: str, donor: Any) -> None:
        """Imports json object from donor to master."""
        for m_dict, d_dict in (
            (self.learn_a_song, donor.learn_a_song),
            (self.play_next, donor.play_next),
            (self.score_attack, donor.score_attack),
            (self.stats, donor.stats)
        ):
            if arr_id in d_dict:
                m_dict[arr_id] = d_dict[arr_id]

def main() -> None:
    """Provide basic command line interface to song list creator."""
    parser = argparse.ArgumentParser(
        description="Minimal command line for profile merge. Code assumes target "
        "profiles are in the steam folders and will be chosen interactively "
        "profiles in steam folders. Code deletes any working sets in the "
        "working directory."
    )
    parser.add_argument(
        "working_dir",
        help="Working directory for database, config files, working "
        "sub-directories amd working files.",
    )

    args = parser.parse_args()

    profile_mgr = RSProfileManager(
        base_dir=Path(args.working_dir).resolve(True),
        auto_setup=True,
        flush_working_set=True,
    )

    master_id = profile_mgr.cl_choose_profile(
        no_action_text="Exit.",
        header_text=(
            "Choose the target profile. Donor file arrangements that meet the merge "
            "criterion will"
            "\n be copied into this profile from the donor file)."
        ),
    )

    if not master_id:
        return

    donor_id = profile_mgr.cl_choose_profile(
        no_action_text="Exit.",
        header_text=(
            "Choose the donor profile. This file will not be changed, but will provide "
            "data for merging."
        ),
    )

    if not donor_id:
        return

    if master_id == donor_id:
        print("One profile selected as master and donor")
        return

    # get the data
    donor = ProfileArrData(profile_mgr.get_json_subtree(donor_id, []))
    master = ProfileArrData(profile_mgr.get_json_subtree(master_id, []))

    # do the merge - at some point ubi updated format and played count is now in stats.
    # take this as definitive.
    for arr_id, song in donor.stats.items():
        d_count = song.get("PlayedCount", 0)
        if arr_id in master.stats:
            if d_count > master.stats[arr_id].get("PlayedCount", 0):
                # donor either has a play count greater than master, or has a record
                # where master doesn't. Do the import across the board.
                master.transfer_arr(arr_id, donor)

    profile_mgr.mark_as_dirty(master_id)
    profile_mgr.write_files()
    profile_mgr.move_updates_to_steam(profile_mgr.steam_account_id)

if __name__ == "__main__":
    main()