#!/usr/bin/python3
"""Licensed Xbox launches through Xodus and the game's selected ARM64 Proton."""
import hmac
import json
import os
from pathlib import Path
import runpy
import secrets
import shutil
import socket
import subprocess
import sys
import tempfile
import threading

BIN = Path(__file__).resolve().parent
RUNTIME = Path('/usr/local/lib/droiddeck-xbox')
DLLS = ('xgameruntime.dll', 'winhttp.dll', 'windows.web.dll', 'twinapi.appcore.dll', 'wintypes.dll', 'libnghttp2.dll')
CLASSES = {
    'twinapi.appcore.dll': ('Windows.ApplicationModel.DataTransfer.DataTransferManager',
                            'Windows.Security.ExchangeActiveSyncProvisioning.EasClientDeviceInformation',
                            'Windows.System.Profile.AnalyticsInfo', 'Windows.System.UserProfile.AdvertisingManager',
                            'Windows.UI.ViewManagement.ApplicationView'),
    'windows.web.dll': ('Windows.Data.Json.JsonArray', 'Windows.Data.Json.JsonObject', 'Windows.Data.Json.JsonValue'),
    'wintypes.dll': ('Windows.Foundation.PropertyValue', 'Windows.Foundation.Collections.PropertySet',
                     'Windows.Foundation.Metadata.ApiInformation', 'Windows.Storage.Streams.Buffer',
                     'Windows.Storage.Streams.DataWriter'),
}
CLASS_KEY = 'Software\\Microsoft\\WindowsRuntime\\ActivatableClassId\\'
VIEW_DRIVES = 'xwvutsrqponmlkjihgfe'


def bridge(upstream):
    key = secrets.token_hex(32).encode('ascii')
    server = socket.socket()
    server.bind(('127.0.0.1', 0))
    server.listen(16)
    clients = set()
    lock = threading.Lock()
    stopped = threading.Event()

    def forward(source, target):
        try:
            while not stopped.is_set():
                data = source.recv(65536)
                if not data:
                    break
                target.sendall(data)
        except OSError:
            pass
        finally:
            try:
                target.shutdown(socket.SHUT_WR)
            except OSError:
                pass

    def connect(client):
        remote = None
        try:
            client.settimeout(10)
            supplied = b''
            while len(supplied) < len(key):
                part = client.recv(len(key) - len(supplied))
                if not part:
                    return
                supplied += part
            if not hmac.compare_digest(supplied, key):
                return
            remote = socket.socket(socket.AF_UNIX)
            remote.connect(upstream)
            client.settimeout(None)
            with lock:
                clients.add(remote)
            writer = threading.Thread(target=forward, args=(client, remote), daemon=True)
            writer.start()
            forward(remote, client)
            writer.join(2)
        except OSError:
            pass
        finally:
            for sock in (client, remote):
                if sock is not None:
                    with lock:
                        clients.discard(sock)
                    sock.close()

    def accept():
        while not stopped.is_set():
            try:
                client, _ = server.accept()
            except OSError:
                break
            with lock:
                if len(clients) >= 64:
                    client.close()
                    continue
                clients.add(client)
            threading.Thread(target=connect, args=(client,), daemon=True).start()

    threading.Thread(target=accept, daemon=True).start()

    def close():
        stopped.set()
        server.close()
        with lock:
            for client in clients:
                try:
                    client.shutdown(socket.SHUT_RDWR)
                except OSError:
                    pass
                client.close()
            clients.clear()

    return server.getsockname()[1], key.decode('ascii'), close


def register_classes(proton, prefix, env):
    try:
        registry = (prefix / 'system.reg').read_text(encoding='utf-8', errors='replace')
    except OSError:
        return
    missing = [(name, dll) for dll, names in CLASSES.items() for name in names
               if '[%s]' % (CLASS_KEY + name).replace('\\', '\\\\') not in registry]
    if not missing:
        return
    text = 'Windows Registry Editor Version 5.00\r\n\r\n' + ''.join(
        '[HKEY_LOCAL_MACHINE\\%s%s]\r\n"DllPath"="C:\\\\windows\\\\system32\\\\%s"\r\n\r\n' % (CLASS_KEY, name, dll)
        for name, dll in missing)
    with tempfile.TemporaryDirectory(prefix='droiddeck-xbox-') as temporary:
        source = Path(temporary) / 'classes.reg'
        source.write_bytes(b'\xff\xfe' + text.encode('utf-16-le'))
        try:
            failed = subprocess.run([proton, 'run', 'reg.exe', 'import', 'Z:' + str(source).replace('/', '\\')],
                                    stdout=sys.stderr, env=env, timeout=90).returncode != 0
        except subprocess.TimeoutExpired:
            failed = True
    if failed:
        print('droiddeck-xbox: Windows Runtime class registration failed; launching without it', file=sys.stderr)


def view_drive(prefix, exe):
    path = Path(exe)
    view = next((parent for parent in path.parents if parent.name.startswith('droiddeck-xbox-')), None)
    if not path.is_absolute() or view is None:
        return exe, None
    devices = prefix / 'dosdevices'
    for letter in VIEW_DRIVES:
        link = devices / (letter + ':')
        staged = devices / ('.%s-%s' % (letter, secrets.token_hex(4)))
        try:
            if os.path.lexists(link) and (os.path.exists(link) or not os.readlink(link).startswith(str(view.parent.parent) + '/')):
                continue
            devices.mkdir(parents=True, exist_ok=True)
            staged.symlink_to(view)
            os.replace(staged, link)
        except OSError:
            if os.path.lexists(staged):
                staged.unlink()
            continue
        return letter.upper() + ':\\' + str(path.relative_to(view)).replace('/', '\\'), link
    return exe, None


def wine(argv, source_env=None):
    source_env = os.environ if source_env is None else source_env
    with open(source_env['DROIDDECK_XBOX_COMMAND_FILE'], encoding='utf-8') as file:
        command = json.load(file)
    if len(command) < 3 or command[1] != 'waitforexitandrun' or len(argv) != 1:
        raise ValueError('Invalid Xbox Proton invocation')
    env = dict(source_env)
    for name in ('DROIDDECK_XBOX_CREDENTIAL_SOCKET', 'DROIDDECK_XBOX_CREDENTIAL_PORT', 'DROIDDECK_XBOX_CREDENTIAL_KEY'):
        env.pop(name, None)
    env.pop('DROIDDECK_XBOX_COMMAND_FILE', None)
    prefix = Path(source_env['STEAM_COMPAT_DATA_PATH']) / 'pfx'
    prefix.parent.mkdir(parents=True, exist_ok=True)
    if not (prefix / 'system.reg').is_file():
        result = subprocess.run([command[0], 'run', 'cmd', '/c', 'exit', '0'], stdout=sys.stderr, env=env)
        if result.returncode:
            return result.returncode
    system = prefix / 'drive_c/windows/system32'
    system.mkdir(parents=True, exist_ok=True)
    for name in DLLS:
        target = system / name
        if target.is_symlink():
            target.unlink()
        shutil.copyfile(RUNTIME / 'windows' / name, target)
    register_classes(command[0], prefix, env)
    overrides = ';'.join(name[:-4] + '=n' for name in DLLS if name != 'libnghttp2.dll')
    env['WINEDLLOVERRIDES'] = overrides + ';' + env.get('WINEDLLOVERRIDES', '')
    exe, drive = view_drive(prefix, argv[0])
    try:
        return subprocess.call([*command[:2], exe, *command[3:]], env=env)
    finally:
        if drive is not None:
            try:
                drive.unlink()
            except OSError:
                pass


def renew_lease(ask, channel, product, lease, stopped):
    while not stopped.wait(15):
        response = ask(channel, {'op': 'xbox-renew', 'id': product, 'lease': lease}, timeout=10)
        if response is not None and response.get('ok') is False:
            return


def launch(command, env):
    if len(command) < 3 or command[1] != 'waitforexitandrun' or env.get('DROIDDECK_XBOX_ACTIVE') == '1':
        return None
    try:
        store = runpy.run_path(str(BIN / 'droiddeck-store-launch'))
    except Exception as error:
        print('droiddeck-xbox: store detection unavailable (%s); stock launch' % error, file=sys.stderr)
        return None
    found = store['find_store_game'](command[2:])
    if found is None or found[1].get('store') != 'xbox':
        return None
    folder, sidecar = found
    for name in DLLS:
        if not (RUNTIME / 'windows' / name).is_file():
            raise RuntimeError('Xbox Windows runtime is missing; update the Linux runtime')
    launch_dir = env.get('BL_LAUNCH_DIR')
    if not launch_dir:
        raise RuntimeError('No Xbox account channel; start the game from DroidDeck')
    reply = store['ask'](Path(launch_dir) / 'stores', {'op': 'xbox-prepare', 'id': sidecar['id']}, timeout=190)
    if not reply or not reply.get('ok'):
        raise RuntimeError('Xbox account service is unavailable; sign in again and retry')
    stopped = threading.Event()
    close = lambda: None
    try:
        port, key, close = bridge(reply['socket'])

        threading.Thread(target=renew_lease, args=(store['ask'], Path(launch_dir) / 'stores',
                         sidecar['id'], reply['lease'], stopped), daemon=True).start()
        with tempfile.TemporaryDirectory(prefix='droiddeck-xbox-') as temporary:
            context = Path(temporary) / 'command.json'
            context.write_text(json.dumps(command), encoding='utf-8')
            context.chmod(0o600)
            owned_env = dict(env, DROIDDECK_XBOX_ACTIVE='1', DROIDDECK_XBOX_COMMAND_FILE=str(context),
                             XODUS_SOCKET=reply['socket'], XODUS_BRIDGE_PORT=str(port), XODUS_BRIDGE_KEY=key,
                             XODUS_LOG='off', RUST_LOG='off', XODUS_SAVE_SYNC='auto' if sidecar.get('cloud', True) else 'off')
            for name in ('DROIDDECK_XBOX_CREDENTIAL_SOCKET', 'DROIDDECK_XBOX_CREDENTIAL_PORT', 'DROIDDECK_XBOX_CREDENTIAL_KEY'):
                owned_env.pop(name, None)
            if reply.get('playfab'):
                owned_env['XODUS_XSTS_PLAYFAB'] = reply['playfab']
            content = folder / '.xbox-content'
            roots = []
            if content.is_dir():
                for entry in content.iterdir():
                    try:
                        marker = json.loads((entry / '.droiddeck-content.json').read_text())
                        if marker.get('parent') == sidecar['id'] and marker.get('productId') == entry.name and not entry.is_symlink():
                            roots.append('Z:' + str(entry).replace('/', '\\'))
                    except (OSError, ValueError, TypeError):
                        pass
            if roots:
                owned_env['XODUS_CONTENT_ROOTS'] = ';'.join(roots)
            return wine([reply['exe']], owned_env)
    finally:
        stopped.set()
        close()
        store['ask'](Path(launch_dir) / 'stores', {'op': 'xbox-released', 'lease': reply['lease']})


if __name__ == '__main__':
    try:
        sys.exit(wine(sys.argv[1:]))
    except (OSError, ValueError, KeyError, RuntimeError):
        print('droiddeck-xbox: Xbox launch failed; check your account, runtime and package.', file=sys.stderr)
        sys.exit(1)
