#!/usr/bin/python3
# SPDX-License-Identifier: LicenseRef-Proprietary
"""Save an explicit default desktop and restore it once per Linux boot."""
import argparse
from collections import Counter
import fcntl
import json
import os
from pathlib import Path
import sys
import tempfile
import time

SERVICE = 'org.lyku.Gravitas.Session1'
OBJECT = '/org/lyku/Gravitas/Session1'
EXCLUDED = {'org.lyku.Gravitas.Settings'}


def paths():
    config = Path(os.environ.get('XDG_CONFIG_HOME') or Path.home() / '.config') / 'gravitas'
    state = Path(os.environ.get('XDG_STATE_HOME') or Path.home() / '.local/state') / 'gravitas'
    return config / 'startup-layout.json', state / 'layout-boot'


def atomic_write(path, text):
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    fd, name = tempfile.mkstemp(prefix='.' + path.name, dir=path.parent)
    try:
        with os.fdopen(fd, 'w') as stream:
            stream.write(text)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(name, path)
    finally:
        if os.path.exists(name):
            os.unlink(name)


def read(path):
    if not path.exists():
        return {'version': 1, 'enabled': False, 'windows': [], 'skipped': []}
    if path.stat().st_size > 262144:
        raise ValueError('Saved layout is too large')
    data = json.loads(path.read_text())
    if data.get('version') != 1 or type(data.get('enabled')) is not bool:
        raise ValueError('Unsupported saved layout')
    windows = data.get('windows')
    if not isinstance(windows, list) or len(windows) > 128:
        raise ValueError('Invalid saved window list')
    for entry in windows:
        if not isinstance(entry, dict) or not isinstance(entry.get('desktop_id'), str):
            raise ValueError('Invalid saved application')
        ident = entry['desktop_id']
        if '/' in ident or not ident.endswith('.desktop') or len(ident) > 1024:
            raise ValueError('Invalid desktop entry ID')
        p = entry.get('placement', {})
        for key in ('app_id', 'output'):
            if not isinstance(p.get(key), str) or not p[key] or len(p[key]) > 1024:
                raise ValueError('Invalid saved window identity')
        for key, low, high in [('x', -131072, 131072), ('y', -131072, 131072),
                               ('width', 1, 32768), ('height', 1, 32768)]:
            if type(p.get(key)) is not int or not low <= p[key] <= high:
                raise ValueError('Invalid saved window geometry')
        if any(type(p.get(key)) is not bool for key in ('maximized', 'fullscreen')):
            raise ValueError('Invalid saved window state')
    return data


def catalog():
    from gi.repository import Gio
    return [app for app in Gio.AppInfo.get_all()
            if isinstance(app, Gio.DesktopAppInfo) and not app.get_is_hidden()]


def resolve(app_id, apps):
    # Prefer exact desktop IDs; WM_CLASS is only usable when it identifies one app.
    exact = [a for a in apps if a.get_id().removesuffix('.desktop') == app_id.removesuffix('.desktop')]
    if len(exact) == 1:
        return exact[0]
    matches = [a for a in apps if (a.get_startup_wm_class() or '').casefold() == app_id.casefold()
               or a.get_id().removesuffix('.desktop').casefold() == app_id.removesuffix('.desktop').casefold()]
    return matches[0] if len(matches) == 1 else None


def capture(windows, apps):
    entries, skipped = [], set()
    for placement in windows:
        ident = placement['app_id']
        if ident in EXCLUDED:
            continue
        app = resolve(ident, apps)
        if app is None:
            skipped.add(ident)
            continue
        entries.append({'desktop_id': app.get_id(), 'name': app.get_display_name(), 'placement': placement})
    if not entries:
        raise ValueError('No restorable application windows are open. Open your apps, arrange them, then save again.')
    return {'version': 1, 'enabled': True, 'windows': entries, 'skipped': sorted(skipped)}


def call(method, placements=None):
    from gi.repository import Gio, GLib
    bus = Gio.bus_get_sync(Gio.BusType.SESSION, None)
    args = None if placements is None else GLib.Variant('(s)', (json.dumps(placements),))
    result = bus.call_sync(SERVICE, OBJECT, SERVICE, method, args,
                           GLib.VariantType('(s)'), Gio.DBusCallFlags.NONE, 5000, None)
    return json.loads(result.unpack()[0])


def claim_boot(marker, boot):
    if marker.exists() and marker.read_text() == boot:
        return False
    # Claim before launching: a helper/compositor crash must not spawn another batch.
    atomic_write(marker, boot)
    return True


def restore(data, marker, boot):
    if not data['enabled'] or not data['windows']:
        return
    if marker.exists() and marker.read_text() == boot:
        return
    apps = {a.get_id(): a for a in catalog()}
    pending = call('BeginLayoutRestore', [w['placement'] for w in data['windows']])
    claim_boot(marker, boot)
    failures = []
    try:
        needed = Counter(p['app_id'] for p in pending)
        for entry in data['windows']:
            ident = entry['placement']['app_id']
            if not needed[ident]:
                continue
            needed[ident] -= 1
            app = apps.get(entry['desktop_id'])
            if app is None:
                failures.append(entry['desktop_id'] + ': no longer installed')
                continue
            try:
                if not app.launch([], None):
                    raise RuntimeError('application refused to launch')
            except Exception as error:
                failures.append(entry['desktop_id'] + ': ' + str(error))
        deadline = time.monotonic() + 45
        while pending and time.monotonic() < deadline:
            time.sleep(0.25)
            pending = call('PendingLayoutRestore')
        if pending:
            failures.append('Windows not opened: ' + ', '.join(p['app_id'] for p in pending))
    finally:
        # Expire unmatched entries so a manually opened window later cannot jump.
        call('BeginLayoutRestore', [])
    if failures:
        raise RuntimeError('; '.join(failures))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['status', 'save', 'enable', 'disable', 'clear', 'restore'])
    action = parser.parse_args().action
    path, marker = paths()
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    with open(path.with_suffix('.lock'), 'a') as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise RuntimeError('A layout operation is already running. Try again when startup has finished.')
        # Replacing or clearing a broken file must remain possible from Settings.
        data = read(path) if action not in ('save', 'clear') else None
        boot = Path('/proc/sys/kernel/random/boot_id').read_text()
        if action == 'restore':
            restore(data, marker, boot)
            return
        if action == 'save':
            data = capture(call('GetWindowLayout'), catalog())
            atomic_write(marker, boot)
        elif action in ('enable', 'disable'):
            if action == 'enable' and not data['windows']:
                raise ValueError('Save a layout first')
            data['enabled'] = action == 'enable'
            atomic_write(marker, boot)
        elif action == 'clear':
            data = {'version': 1, 'enabled': False, 'windows': [], 'skipped': []}
        if action != 'status':
            atomic_write(path, json.dumps(data, indent=2) + '\n')
        print(json.dumps(data))


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        print('Startup layout: ' + str(error), file=sys.stderr)
        sys.exit(1)
