#!/usr/bin/python3 -s
""" Script that automatically installs protonfixes in all Proton installations
"""

import glob
import os
import re
import shutil
import subprocess
import sys


STEAM_DIRS = ["~/.local/share/Steam", "~/.steam/steam"]
CT_DIRS = ["~/.local/share/Steam/compatibilitytools.d",
           "~/.steam/root/compatibilitytools.d",
           "~/.steam/compatibilitytools.d"]
SYS_CT_DIRS = ["/usr/share/steam/compatibilitytools.d",
               "/usr/local/share/steam/compatibilitytools.d"]
LF_PATH = 'steamapps/libraryfolders.vdf'
PROTON_PATH = 'steamapps/common/Proton'

LIBRARY_RE = re.compile(r'"\d*"\s*"(?P<path>(\/\w*)*)"')
CTOOLS_RE = re.compile(r'"install_path"\s*"(?P<path>(\/\w*)*)"')


def find_steam_installation():
    """ Finds the steam user installation folder by looking in STEAM_DIRS
    """
    for path in STEAM_DIRS:
        dirpath = os.path.expanduser(path)
        if os.path.isdir(dirpath):
            return dirpath
    raise NotADirectoryError("Steam directory not found")


def find_libraries(steam_dir):
    """ Finds additional library folders with the LIBRARY_RE regex
    """
    lf_path = os.path.join(steam_dir, LF_PATH)
    return find_regex_groups(lf_path, LIBRARY_RE, 'path')


def find_protons(library_dirs):
    """ Finds the proton directories given a list of library directories
    """
    dirs = []
    for directory in library_dirs:
        preglob_proton = os.path.join(directory, PROTON_PATH)
        for proton_dir in glob.glob(preglob_proton + '*'):
            dirs.append(proton_dir)
    return dirs


def scan_compatibilitytools(dirlist):
    """ Scans a list of directories and finds proton instances within the
        compatibilitytools.d standard
    """
    temp_dirs = []
    dirlist_f = [x for x in dirlist
                 if os.path.isdir(os.path.expanduser(x))]
    for directory in dirlist_f:
        ex_dir = os.path.expanduser(directory)
        for entry in os.listdir(ex_dir):
            path = os.path.join(ex_dir, entry)
            if os.path.isdir(path):
                temp_dirs.append(path)
            elif os.path.isfile(path) and path.endswith('vdf'):
                temp_dirs += find_regex_groups(path, CTOOLS_RE, 'path')
    return [x for x in temp_dirs if check_proton(x)]


def find_soldier(library_dirs):
    """ Find soldier runtime installation given a list of library directiries
    """
    for directory in library_dirs:
        test_path = os.path.join(directory, 'steamapps/common/SteamLinuxRuntime_soldier')
        if os.path.isdir(test_path):
            return test_path
    return None


def check_proton(path):
    """ Checks for the presence of common proton files
    """
    return (os.path.isdir(path)
            and checkfile(path, 'proton')
            and (checkfile(path, 'user_settings.sample.py')
                 or checkfile(path, 'user_settings.py'))
            and checkfile(path, 'toolmanifest.vdf')
            and not (checkfile(path, 'protonfixes')  # Weed out protons with bundled protonfixes
                     or os.path.isdir(os.path.join(path, 'protonfixes')))
            )


def checkfile(path, filename):
    """ Checks if file filename exists in path
    """
    return os.path.isfile(os.path.join(path, filename))


def find_regex_groups(path, regex, groupname):
    """ Given a file and a regex with a named group groupname, return an
        array of all the matches
    """
    matches = []
    with open(path) as re_file:
        for line in re_file:
            search = regex.search(line)
            if search:
                matches.append(search.group(groupname))
    return matches


def install_protonfixes(proton_dir):
    """ Installs protonfixes given the specified proton directory
    """
    us_file = os.path.join(proton_dir, 'user_settings.py')
    if not os.path.isfile(us_file):
        us_ex_file = os.path.join(proton_dir, 'user_settings.sample.py')
        shutil.copy(us_ex_file, us_file)

    protonfixes_installed = False
    with open(us_file, 'r') as user_settings:
        for line in user_settings:
            if line.startswith("import protonfixes"):
                print("Protonfixes already installed in {}, skipping".format(proton_dir))
                protonfixes_installed = True
    if not protonfixes_installed:
        print("Installing protonfixes in {}".format(proton_dir))
        with open(us_file, 'a') as user_settings:
            user_settings.write("\nimport protonfixes")


def find_protonfixes():
    """ Find the protonfixes directory from PYTHONPATH
    """
    for path_dir in sys.path:
        candidate = os.path.join(path_dir, 'protonfixes')
        if (os.path.isdir(candidate)
            and os.path.isfile(os.path.join(candidate, '__init__.py'))):
            return candidate
    return None


def patch_soldier(soldier_dir):
    """ Given the soldier directory launch soldier_protonfixify with the
        right arguments
    """
    protonfixes_dir = find_protonfixes()
    subprocess.run(("soldier_protonfixify", soldier_dir, protonfixes_dir), check=False)


def main():
    """ Main function
    """
    try:
        proton_dirs = []
        soldier_dir = None
        if os.getuid() != 0:
            steam_dir = find_steam_installation()
            library_dirs = [steam_dir, ] + find_libraries(steam_dir)
            proton_dirs += find_protons(library_dirs)
            proton_dirs += scan_compatibilitytools(CT_DIRS + SYS_CT_DIRS)
            soldier_dir = find_soldier(library_dirs)
        else:
            print("install_protonfixes run as root, only scanning system compatibilitytools.d")
            proton_dirs += scan_compatibilitytools(SYS_CT_DIRS)
        proton_dirs_final = sorted({os.path.realpath(x) for x in proton_dirs})
        for pdir in proton_dirs_final:
            try:
                install_protonfixes(pdir)
            except IOError:
                print("Protonfixes installation failed in {}, skipping".format(pdir))
        if soldier_dir:
            print("Soldier runtime detected, patching it")
            patch_soldier(soldier_dir)
    except NotADirectoryError:
        print("Could not find steam installation directory")


if __name__ == "__main__":
    main()
