#!/usr/bin/env python
from steamctl.clients import CachingSteamClient
import sys
import json
import requests
from pathlib import Path
from argparse import Namespace
from io import StringIO
import winreg

steamClient = False
userdata = None

class App:
    """
    A class representing an app added to Steam.
    
    This information is pulled from the shortcuts.vdf file with the
    #parseShortcuts method. 
    
    Attributes
    ----------
    appId : int
        The ID of the application as represented in steam
    name : str
        The display name of the app
    executable : str
        The link to the executable file.
    start : str, Path
        The path to start the app in
    icon : str
        The path to the icon file that Steam uses. If empty, then the icon of the executable is used in Windows installations.
    shortcut : str
    options : str
        Command-line arguments to use when launching the executable.
    hidden : bool
        Whether or not the app is hidden from the client. Defaults to false.
    desktopConfig : bool
    overlay : bool
        Whether or not the steam overlay can display on this game.
    vr : bool
        Whether or not the app is VR enabled.
    devkit : bool
    devkitId: str
    playTime: long
        The date that the application was last launched.
    tags: str
    """
    
    appId = 0
    name = ""
    executable = ""
    start = ""
    icon = ""
    shortcut = ""
    options = ""
    hidden = False
    desktopConfig = False
    overlay = False
    vr = False
    devkit = False
    devkitId = ""
    playTime = 0
    tags = "null"
    
    def __init__(self, appId):
        self.appId = appId
    
    def __repr__(self):
        return "App %s (%s at %s)" % (self.appId, self.name, self.executable)
    
    def __hash__(self):
        return self.appId
    
    def __eq__(self, other):
        if self.appId != other.appId:
            return False
        if self.name != other.name:
            return False
        if self.executable != other.executable:
            return False
        if self.start != other.start:
            return False
        if self.icon != other.icon:
            return False
        if self.shortcut != other.shortcut:
            return False
        if self.options != other.options:
            return False
        if self.hidden != other.hidden:
            return False
        if self.desktopConfig != other.desktopConfig:
            return False
        if self.overlay != other.overlay:
            return False
        if self.vr != other.vr:
            return False
        if self.devkit != other.devkit:
            return False
        if self.devkitId != other.devkitId:
            return False
        if self.playTime != other.playTime:
            return False
        if self.tags != other.tags:
            return False
        return True

def parseShortcuts(shortcutFile):
    """Parses the "userdata/<userId>/config/shortcuts.vdf" file.
    
    This file contains all the information about what non-steam apps are
    installed in the steam client. This method parses the bytecode and returns
    the apps obtained.
    
    Parameters
    ----------
    shortcutFile : str, Path
        The path to the shortcuts.vdf file to parse.
       
    Returns
    -------
    list
        A list of apps contained within the shortcut file.
    """
    
    data = []
    with open(shortcutFile, 'rb') as shortcutData:
        rawData = shortcutData.read()
        i = 0
        buffer = bytearray()
        app = None
        while i < len(rawData):
            buffer.append(rawData[i])
            try:
                buffCont = buffer.decode()
            except UnicodeDecodeError:
                #print(rawData[:i])
                #print("Clearing buffer %s after %d" % (buffer[:-1].decode(), buffer[-1]))
                buffer.clear()
                buffCont = ""
            found = True
            if buffCont.endswith("shortcuts"):
                buffer.clear()
                i += 2
                while rawData[i] != 2:
                    buffer.append(rawData[i])
                    i += 1
                data.append({
                    "shortcuts": buffer[:-1].decode()
                })
            if buffCont.endswith("appid"):
                buffer.clear()
                i += 2
                # Some app names do have a 1 for one of their bits
                while rawData[i + 1:i + 8] != bytearray("AppName", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app = App(int.from_bytes(buffer, "little"))
                data.append(app)
            elif buffCont.endswith("AppName"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 4] != bytearray("Exe", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.name = buffer[:-1].decode()
            elif buffCont.endswith("Exe"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 9] != bytearray("StartDir", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.executable = buffer[:-1].decode()
            elif buffCont.endswith("StartDir"):
                buffer.clear()
                i += 2
                while rawData[i] != 0:
                    buffer.append(rawData[i])
                    i += 1
                app.start = buffer.decode()
            elif buffCont.endswith("icon"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 13] != bytearray("ShortcutPath", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.icon = buffer[:-1].decode()
            elif buffCont.endswith("ShortcutPath"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 14] != bytearray("LaunchOptions", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.shortcut = buffer[:-1].decode()
            elif buffCont.endswith("LaunchOptions"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 9] != bytearray("IsHidden", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.options = buffer[:-1].decode()
            elif buffCont.endswith("IsHidden"):
                buffer.clear()
                i += 2
                while len(buffer) < 3:
                    buffer.append(rawData[i])
                    i += 1
                app.hidden = int.from_bytes(buffer, "little")
            elif buffCont.endswith("AllowDesktopConfig"):
                buffer.clear()
                i += 2
                while len(buffer) < 3:
                    buffer.append(rawData[i])
                    i += 1
                app.desktopConfig = int.from_bytes(buffer, "little")
            elif buffCont.endswith("AllowOverlay"):
                buffer.clear()
                i += 2
                while len(buffer) < 3:
                    buffer.append(rawData[i])
                    i += 1
                app.overlay = int.from_bytes(buffer, "little")
            elif buffCont.endswith("OpenVR"):
                buffer.clear()
                i += 2
                while len(buffer) < 3:
                    buffer.append(rawData[i])
                    i += 1
                app.vr = int.from_bytes(buffer, "little")
            elif buffCont.endswith("Devkit") and type(app.devkit) == bool:
                buffer.clear()
                i += 2
                while len(buffer) < 3:
                    buffer.append(rawData[i])
                    i += 1
                app.devkit = int.from_bytes(buffer, "little")
            elif buffCont.endswith("DevkitGameID"):
                buffer.clear()
                i += 2
                while rawData[i + 1:i + 13] != bytearray("LastPlayTime", "utf-8"):
                    buffer.append(rawData[i])
                    i += 1
                app.devkitId = buffer[:-1].decode()
            elif buffCont.endswith("LastPlayTime"):
                buffer.clear()
                i += 2
                while len(buffer) < 4:
                    buffer.append(rawData[i])
                    i += 1
                app.playTime = int.from_bytes(buffer, "little")
            elif buffCont.endswith("tags"):
                buffer.clear()
                i += 2
                while i < len(rawData) and rawData[i] != 2:
                    buffer.append(rawData[i])
                    i += 1
                if i < len(rawData) and rawData[i] == 2:
                    buffer = buffer[:-1]
                app.tags = buffer.decode()
            else:
                found = False
            #print('\r' + buffCont, end='')
            if found:
                #try:
                    #print(": %s" % buffer.decode())
                #except UnicodeDecodeError:
                    #print(": %s" % int.from_bytes(buffer, "little"))
                buffer.clear()
            i += 1
    return data

def writeShortcuts(apps, shortcutFile):
    """Writes app data to the shortcut file.
    
    The shortcut file will be completely overwritten. As such, we can take the
    opporunity to add whatever apps we want.
    
    Parameters
    ----------
    apps : list
        A list of apps to write.
    shortcutFile : str, Path
        The path to the shortcuts.vdf file to write. It will be created if
        needed.
    """
    
    encoding = "utf-8"
    data = bytearray([0x00])
    for app in apps:
        if type(app) == dict:
            data += bytearray("shortcuts", encoding)
            data.append(0x00)
            data += bytearray(apps[0]["shortcuts"], encoding)
        else:
            data.append(0x00)
            data.append(0x02)
            data += bytearray("appid", encoding)
            data.append(0x00)
            data += int(app.appId).to_bytes(4, "little")
            data.append(0x01)
            data += bytearray("AppName", encoding)
            data.append(0x00)
            data += bytearray(app.name, encoding)
            data.append(0x00)
            data.append(0x01)
            data += bytearray("Exe", encoding)
            data.append(0x00)
            data += bytearray(app.executable, encoding)
            data.append(0x00)
            data.append(0x01)
            data += bytearray("StartDir", encoding)
            data.append(0x00)
            data += bytearray(app.start, encoding)
            data.append(0x00)
            data.append(0x01)
            data += bytearray("icon", encoding)
            data.append(0x00)
            data += bytearray(app.icon, encoding)
            data.append(0x00)
            data.append(0x01)
            data += bytearray("ShortcutPath", encoding)
            data.append(0x00)
            data += bytearray(app.shortcut, encoding)
            data.append(0x00)
            data.append(0x01)
            data += bytearray("LaunchOptions", encoding)
            data.append(0x00)
            data += bytearray(app.options, encoding)
            data.append(0x00)
            data.append(0x02)
            data += bytearray("IsHidden", encoding)
            data.append(0x00)
            data += int(app.hidden).to_bytes(3, "little")
            data.append(0x00)
            data.append(0x02)
            data += bytearray("AllowDesktopConfig", encoding)
            data.append(0x00)
            data += int(app.desktopConfig).to_bytes(3, "little")
            data.append(0x00)
            data.append(0x02)
            data += bytearray("AllowOverlay", encoding)
            data.append(0x00)
            data += int(app.overlay).to_bytes(3, "little")
            data.append(0x00)
            data.append(0x02)
            data += bytearray("OpenVR", encoding)
            data.append(0x00)
            data += int(app.vr).to_bytes(3, "little")
            data.append(0x00)
            data.append(0x02)
            data += bytearray("Devkit", encoding)
            data.append(0x00)
            data += int(app.devkit).to_bytes(3, "little")
            data.append(0x00)
            data.append(0x01)
            data += bytearray("DevkitGameID", encoding)
            data.append(0x00)
            data += bytearray(app.devkitId, encoding)
            data.append(0x00)
            data.append(0x02)
            data += bytearray("LastPlayTime", encoding)
            data.append(0x00)
            data += int(app.playTime).to_bytes(4, "little")
            data.append(0x00)
            data += bytearray("tags", encoding)
            data.append(0x00)
            data += bytearray(app.tags, encoding)
    with open(shortcutFile, 'wb') as stream:
        stream.write(data)

def getSteamArgs():
    """Logs into the steam client as an anonymous user. This allows us to
    download information on apps not installed.
    
    Returns
    -------
    The interface to Steam.
    """
    
    global steamClient
    if not steamClient:
        args = Namespace()
        args.user = None
        args.anonymous = True
        args.skip_licenses = True
        steamClient = CachingSteamClient()
        print("Logging into steam")
        steamClient.login_from_args(args)
    return steamClient

def getHeader(code):
    """Obtains the URL to the header image of a game.
    
    Parameters
    ----------
    code : int, str
        The app ID of the game to get the header for.
    
    Returns
    -------
    str
        The URL to the game's header image.
    """
    
    return "https://cdn.cloudflare.steamstatic.com/steam/apps/%s/header.jpg" % code

def getLibraryTall(code):
    """Obtains the URL to the box art image of a game.
    
    Parameters
    ----------
    code : int, str
        The app ID of the game to get the box art for.
    
    Returns
    -------
    str
        The URL to the game's box art image.
    """
    
    return "https://cdn.cloudflare.steamstatic.com/steam/apps/%s/library_600x900.jpg" % code

def getBackground(code):
    """Obtains the URL to the page background image of a game.
    
    Parameters
    ----------
    code : int, str
        The app ID of the game to get the background for.
    
    Returns
    -------
    str
        The URL to the game's page background image.
    """
    
    return "https://cdn.cloudflare.steamstatic.com/steam/apps/%s/library_hero.jpg" % code

def getLogo(code):
    """Obtains the URL to the logo image of a game.
    
    Parameters
    ----------
    code : int, str
        The app ID of the game to get the logo for.
    
    Returns
    -------
    str
        The URL to the game's logo image.
    """
    
    return "https://cdn.cloudflare.steamstatic.com/steam/apps/%s/logo.png" % code

def getGameInfo(code):
    """Deprecated"""
    fromContent = requests.get(fromURL, allow_redirects=True)
    return "https://cdn.cloudflare.steamstatic.com/steam/apps/%s/logo.png" % code

def copyContent(name, content, byte=True):
    """
    Adds a file to all users.
    
    Parameters
    ----------
    name : str
        The name of the file to add.
    content : str, list
        The content to write to the file.
    byte : bool, optional
        Whether or not to write in bytenode. Defaults to true.
    """
    
    for user in userdata.iterdir():
        toFile = Path(user, name)
        if not toFile.parent.is_dir():
            toFile.parent.mkdir(parents=True)
        print("Writing to %s" % (toFile))
        mode = 'w'
        if byte:
            mode += 'b'
        with open(toFile, mode) as stream:
            stream.write(content)

def copyFile(fromURL, toCode):
    """Downloads an image from one game on Steam and registers it under another
    game.
    
    This essentially just downloads the image and stores it under the
    appropriate name for the Steam client to find.
    
    Parameters
    ----------
    fromURL : str, URL
        The URL of the image to download. The name of the file must be
        "header.jpg," "library_600x900.jpg," "library_hero.jpg," or "logo.png."
    toCode : int
        The app ID of the game to register the images to.
    """
    
    fromURLComp = fromURL.split('/')
    fromName = fromURLComp[6]
    fromCode = fromURLComp[5]\
    
    # Gets the final file name
    if fromName == "library_hero.jpg":
        toName = "_hero.jpg"
    elif fromName == "library_600x900.jpg":
        toName = "p.jpg"
    elif fromName == "header.jpg":
        toName = ".jpg"
    else:
        toName = "_" + fromName
    toName = toCode + toName
    
    print("Copying %s (%s) to %s" % (fromName, fromCode, toName))
    fromContent = requests.get(fromURL, allow_redirects=True)
    
    copyContent(Path("config", "grid", toName), fromContent.content)

def copyContents(ids):
    """Goes through a list of app ID mappings, downloading information on
    "from" titles and mapping it to "to" titles, before writting that
    information to the disk for the steam client to find.
    
    Parameters
    ----------
    ids : list
        A two-dimensional list of numbers, where the first index of each
        sub-list contains the app to download information from, and the second
        is the ID of the app to store that information under.
    """
    
    client = getSteamArgs()
    fromCodes = []
    toCodes = []
    reverse = {}
    for id in ids:
        # Get the mappings
        fromCode = id[0]
        toCode = id[1]
        reverse[toCode] = fromCode
        fromCodes.append(fromCode)
        toCodes.append(toCode)
        # Copy the appropriate images from online
        copyFile(getHeader(str(fromCode)), str(toCode))
        copyFile(getLibraryTall(str(fromCode)), str(toCode))
        copyFile(getBackground(str(fromCode)), str(toCode))
        copyFile(getLogo(str(fromCode)), str(toCode))
    # Obtains information on the apps to download from, notably logo positions
    # and such.
    data = client.get_product_info(apps=fromCodes)
    i = 0
    for i in range(len(ids)):
        # Generate the positioning data for logos and such
        logo = data["apps"][ids[i][0]]["common"]["library_assets"]["logo_position"]
        logo2 = {
            "nVersion": 1,
            "logoPosition": {
                "pinnedPosition": logo["pinned_position"],
                "nWidthPct": logo["width_pct"],
                "nHeightPct": logo["height_pct"]
            }
        }
        # Stores the logo positioning to file.
        copyContent(Path("config", "grid", str(ids[i][1]) + ".json"), json.dumps(logo2), byte=False)
        
        # Generate the data for other stuff.
        logo2 = [[
                "achievements", {
                    "version":2,
                    "data": {
                        "vecHighlight": [],
                        "vecUnachieved": [],
                        "vecAchievedHidden": [],
                        "nTotal": 0,
                        "nAchieved": 0
                    }
                }
            ], [
                "customimage", {
                    "version": 1,
                    "data": logo2["logoPosition"]
                }
            ]
        ]
        copyContent(Path("config", "librarycache", str(ids[i][1]) + ".json"), json.dumps(logo2), byte=False)

    # Write the actual shortcut data.
    for user in userdata.iterdir():
        shortcutData = parseShortcuts(Path(user, "config", "shortcuts.vdf"))
        for app in shortcutData:
            if type(app) != dict:
                if app.appId in toCodes:
                    print("Setting shortcut data for %s", app.name)
                    app.name = data["apps"][reverse[app.appId]]["common"]["name"]
                    app.icon = str(Path(user, "config", "grid", str(app.appId) + ".jpg"))
                    print("Setting %s to %s" % (app.name, app.icon))
        writeShortcuts(shortcutData, Path(user, "config", "shortcuts.vdf"))
        
    
if __name__ == "__main__":
    keys = []
    ids = []
    ops = []
    # Gather the parameters
    for term in sys.argv[1:]:
        if not term.startswith('-'):
            # All non-flag arguments are added to keys, whether or not they are
            # a steam game ID.
            keys.append(term)
            # Append the ids
            if term.isnumeric():
                if len(ids) > 0 and type(ids[-1]) == int:
                    ids[-1] = (ids[-1], int(term))
                else:
                    ids.append(int(term))
            else:
                # If we run into an argument that isn't either a flag or a steam
                # id, then we must not have any IDs, and we are searching for a
                # game.
                ids = None
        else:
            if term[1] == '-':
                ops.append(term[2:])
            else:
                for op in term[1:]:
                    ops.append(op)
    if len(keys) < 1 and not "a" in ops:
        print("%s [OPTIONS] (<fromId> <toId>)..." % __file__)
        print("%s [OPTIONS] <term>..." % __file__)
        exit(1)
        
    # Find the steam installation directory.
    try:
        import winreg
        hkey = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, "SOFTWARE\WOW6432Node\Valve\Steam")
        if hkey:
            userdata = Path(winreg.QueryValueEx(hkey, "InstallPath")[0], "userdata")
    except:
        hkey = None
    if not userdata:
        userdata = Path(Path.home(), ".steam", "root", "userdata")
    if not userdata:
        print("Could not find Steam installation directory.", file=sys.stderr)
        exit(1)
        
    if ids:
        # In this event, we make sure that all IDs have a mapping, by checking
        # the last index.
        if type(ids[-1]) == int:
            print("Must have an even number of fromID - toId pairs")
            exit(2)
    
        copyContents(ids)
    else:
        tempApps = {}
        apps = []
        
        # Loop through all user folders
        for user in userdata.iterdir():
            # The shortcut file contain a collection of games that the client recognizes
            shortcutData = parseShortcuts(Path(user, "config", "shortcuts.vdf"))
            # Obtain 
            for app in shortcutData:
                if type(app) != dict:
                    if "a" in ops:
                        # If the all flag was specified, then add all apps to
                        # the list
                        apps.append(app)
                    else:
                        for term in keys:
                            # Otherwise, record how many of the search terms
                            # appear in the app title.
                            if term.lower() in app.name.lower():
                                if not app in tempApps:
                                    tempApps[app] = 1
                                else:
                                    tempApps[app] += 1
        if not "a" in ops:
            # If we aren't listing all apps, find which apps are at least a 70%
            # match.
            for app in tempApps:
                if tempApps[app] / len(keys) >= 0.7:
                    apps.append(app)
        if len(apps) == 0:
            if "v" in ops:
                # Verbose output.
                print("No matches found")
            exit(0)
        elif len(apps) == 1 and not "v" in ops and not "l" in ops:
            # If there is only one match and we aren't looking for details, we
            # just print the app ID and be done with it.
            print(apps[0].appId)
        else:
            for app in apps:
                print("%d - %s" % (app.appId, app.name))
                if "l" in ops:
                    print("\tExecutable: %s" % app.executable)
                    print("\tRuntime Directory: %s" % app.start)
                    print("\tIcon: %s" % app.icon)
                    print("\tShortcut: %s" % app.shortcut)
                    print("\tRuntime Options: %s" % app.options)
                    print("\tIs Hidden: %s" % bool(app.hidden))
                    print("\tDesktop Configuration: %s" % bool(app.desktopConfig))
                    print("\tOverlay Enabled: %s" % bool(app.overlay))
                    print("\tIs VR: %s" % bool(app.vr))
                    print("\tDevkit: %d" % app.devkit)
                    print("\tDevkit ID: %s" % app.devkitId)
                    print("\tDate of last Play: %s" % app.playTime)
                    print("\tTags: %s" % app.tags)