#!/usr/bin/env python3
'''
Run a Steam game through SteamFS.

Copyright (C) 2019,2020       Valve, Inc

This program can be distributed under the terms of the GNU GPL.
'''

import sys
import subprocess
import time
import logging
import contextlib
import os
import datetime
import argparse
import signal

from pathlib import Path
from typing import Sequence

# Maximum time to wait for mountpoint to come up and go down
TIMEOUT = 30

# Arguments that are passed through to SteamFS
PASSTHROUGH_ARGS = [
    ('--ondemand', bool),
    ('--download-blocksize', int),
    ('--cache-blocksize', int),
    ('--cachefile', str),
    ('--readahead-profile', str)
]

log = logging.getLogger(__name__)

def parse_args(argv):
    progname = os.path.basename(argv[0])
    parser = argparse.ArgumentParser(
        description='Run SteamFS and execute given program in the mountpoint.',
        usage=f'{progname} [--ondemand] <executable> -- [args..]')

    parser.add_argument('--no-bpf', action='store_true', default=False,
                        help='Disable BPF-based read logging.')
    parser.add_argument('executable', type=str,
                        help='Executable to run through SteamFS.')
    parser.add_argument('--logdir', type=str, default=os.path.expanduser('~/steamfs_logs/'),
                        help='Where to save read logs (default: ~/steamfs_logs)')
    parser.add_argument('--crashdir', type=str, default=os.path.expanduser('~/steamfs_logs/'),
                        help='Where to save crash logs (default: ~/steamfs_logs)')
    parser.add_argument('args', nargs='*',
                        help='Arguments to pass to executable.')

    for (arg, type_) in PASSTHROUGH_ARGS:
        if type_ == bool:
            parser.add_argument(
                arg, action='store_true', default=None, dest=arg,
                help='(passed through to SteamFS)')
        else:
            parser.add_argument(
                arg, type=type_, default=None, dest=arg,
                help='(passed through to SteamFS)')

    return parser.parse_args(argv[1:])


def get_win_exe(args: Sequence[str]) -> Path:
    '''Extract Windows binary name from Proton argument list.'''

    for (i, arg) in enumerate(args):
        if arg.startswith('-') or not arg.endswith('.exe'):
            continue
        p = Path(arg)
        if p.is_file():
            return p
    raise RuntimeError('Unable to find game path in Proton arguments: %r' % args)


def main(exit_stack, args):
    '''Start the game and return its runtime'''

    for path in (args.crashdir, args.logdir):
        if not os.path.exists(path):
            os.mkdir(path)

    if 'STEAM_COMPAT_INSTALL_PATH' in os.environ:
        mountpoint = Path(os.environ['STEAM_COMPAT_INSTALL_PATH'])
    else:
        # When using Proton, we have to account for the fact that Proton may be installed
        # in a different Steam library - and we want to use SteamFS for the one that
        # contains the game, not Proton.
        if args.executable.endswith('proton'):
            (mountpoint, _) = split_path(get_win_exe(args.args))
        else:
            (mountpoint, _) = split_path(Path(args.executable))
    log.info('Mountpoint is %s', mountpoint)

    cwd = os.getcwd()
    ofd = os.open(args.crashdir + datetime.datetime.now().strftime(
        '%Y-%m-%d_%H-%M-%S.log'), os.O_WRONLY | os.O_CREAT, mode=0o644)
    steamfs_cmdline = [find_program('fs/steamfs'), '--log-reads', '--crashdir',
                       args.crashdir, str(mountpoint), str(mountpoint)]
    subprocess_args = {'cwd': args.logdir}

    readahead_profile = getattr(args, '--readahead-profile')
    if readahead_profile and not args.no_bpf:
        setattr(args, '--readahead-profile', None)
    else:
        readahead_profile = None

    for (arg, type_) in PASSTHROUGH_ARGS:
        val = getattr(args, arg)
        if val is None:
            continue
        if type_ is bool and val:
            steamfs_cmdline.append(arg)
        else:
            steamfs_cmdline.append(arg)
            steamfs_cmdline.append(val)

    if getattr(args, '--ondemand'):
        subprocess_args['stderr'] = ofd
    else:
        subprocess_args['stderr'] = subprocess.STDOUT
        subprocess_args['stdout'] = ofd

    log.debug('Running %s', steamfs_cmdline)
    steamfs = subprocess.Popen(steamfs_cmdline, **subprocess_args)
    exit_stack.callback(unmount, steamfs, mountpoint)

    log.info("Waiting for mountpoint to come up...")
    wait_for_mount(steamfs, mountpoint)

    if not args.no_bpf:
        if readahead_profile:
            log.info("Starting kernel-based readahead helper...")
            bpf_process = subprocess.Popen(
                [ find_program('kernel_read_tracing/steamfs_readahead'),
                  readahead_profile, str(mountpoint) ],
                cwd=args.logdir)
            tracer_pid = wait_for_tracer(
                bpf_process, f'{args.logdir}/steamfs_readahead_{bpf_process.pid}.lck')
        else:
            log.info("Starting kernel-based read logger...")
            bpf_process = subprocess.Popen(
                [ find_program('kernel_read_tracing/steamfs_log_reads'),
                  str(mountpoint) ], cwd=args.logdir)
            tracer_pid = wait_for_tracer(
                bpf_process, f'{args.logdir}/steamfs_log_reads_{bpf_process.pid}.lck')
        exit_stack.callback(bpf_process.wait, TIMEOUT)

    cmdline = [ args.executable ] + args.args
    log.info("Executing %s...", cmdline)
    try:
        t1 = time.time()
        # Change into "current" directory, because this may now be in the
        # mountpoint.
        subprocess.run(cmdline, cwd=cwd)
        runtime = time.time() - t1
        log.info('Total runtime: %.3f seconds', runtime)
    finally:
        # Read logger must be started after mounting SteamFS, but does not w
        # terminates on its own until unmount. Therefore, terminate it manually
        if not args.no_bpf:
            log.info('Terminating kernel tracer...')
            try:
                os.kill(tracer_pid, signal.SIGINT)
            except ProcessLookupError:
                log.warning('Tracer has terminated already!?')

    return runtime


def wait_for_tracer(process: subprocess.Popen, path: str) -> int:
    elapsed = 0
    lckfile = Path(path)
    while not lckfile.exists():
        if process.poll() is not None:
            raise RuntimeError('Tracer terminated unexpectedly.')
        if elapsed > TIMEOUT:
            raise RuntimeError('Tracer did not come up.')
        time.sleep(0.1)
        elapsed += 0.1
    return int(lckfile.read_text())


def wait_for_mount(mount_process: subprocess.Popen,
                   mnt_dir: Path):
    elapsed = 0
    while not mnt_dir.is_mount():
        if mount_process.poll() is not None:
            raise RuntimeError('SteamFS terminated unexpectedly.')
        if elapsed > TIMEOUT:
            raise RuntimeError('Mountpoint did not come up.')
        time.sleep(0.1)
        elapsed += 0.1


def unmount(mount_process: subprocess.Popen,
            path: Path):

    # Some games start other processes that continue to access the mountpoint even after
    # the game has terminated.  Therefore, try a few times.
    log.info("Unmounting SteamFS...")
    elapsed = 0
    while True:
        # SteamFS may have been terminated, e.g. with Ctrl-C
        if not path.is_mount():
            break

        if subprocess.run(['fusermount3', '-u', path]).returncode == 0:
            break

        time.sleep(0.1)
        elapsed += 0.1
        if elapsed > TIMEOUT:
            log.error('Unable to unmount SteamFS - falling back to lazy unmount!')
            subprocess.run(['fusermount3', '-u', '-z', path])
            return

    # Give mount process a little while to terminate.
    elapsed = 0
    while mount_process.poll() is None:
        if elapsed > TIMEOUT:
            log.error('Timeout while waiting for SteamFS to terminate, continuing anyway...')
            return
        time.sleep(0.1)
        elapsed += 0.1


def find_program(relpath: str) -> str:
    '''Try to find *relpath* within build directory

    If we can find the build directory, return *relpath*
    within it. Otherwise, return the basename of *relpath*.
    '''

    basedir = Path(__file__).resolve().parent
    if (basedir / 'build.ninja').is_file():
        return str(basedir / relpath)
    else:
        return os.path.basename(relpath)


def split_path(path: Path):
    """Split *path* into <something>/SteamApps/common/<name> and remainder"""

    fpath = path.resolve()
    gamepath = fpath
    while gamepath != gamepath.parent:
        libpath = gamepath.parent.parent.parent
        if ((libpath / 'SteamApps/common').is_dir() or
            (libpath / 'steamapps/common').is_dir()):
            return (gamepath, fpath.relative_to(gamepath))
        gamepath = gamepath.parent
    raise RuntimeError("Unable to find game path")


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    cwd = os.getcwd()
    args = parse_args(sys.argv)
    try:
        with contextlib.ExitStack() as exit_stack:
            runtime = main(exit_stack, args)
    except:
        log.exception('Error while running game through SteamFS')
        runtime = 0

    if runtime < 2:
        log.warning('Game did not start, or ran for less than 2 seconds. Attempting '
                    'to start without SteamFS')
        subprocess.run([ args.executable ] + args.args, cwd=cwd, check=False)
