#!/usr/bin/python3

"""
This script is outside of qubesd because it relies on systemd to:

- Order this action after the autostart or standard qubes;
- Skip preloading if kernel command line prevents autostart.
"""

import asyncio
import concurrent.futures
import qubesadmin


def get_preload_max(qube) -> int | None:
    """
    Get the ``preload-dispvm-max`` feature as an integer.

    :param qubes.vm.qubes.QubesVM qube: Qube to query the feature from.
    """
    value = qube.features.get("preload-dispvm-max", None)
    return int(value) if value else value


async def main():
    app = qubesadmin.Qubes()
    domains = app.domains
    default_dispvm = getattr(app, "default_dispvm", None)
    global_max = get_preload_max(domains["dom0"])
    appvms = [
        qube
        for qube in domains
        if (
            qube.klass == "AppVM"
            and getattr(qube, "template_for_dispvms", False)
            and (
                (qube != default_dispvm and get_preload_max(qube))
                or (
                    (qube == default_dispvm and global_max)
                    or (global_max is None and get_preload_max(qube))
                )
            )
        )
    ]
    method = "admin.vm.CreateDisposable"
    loop = asyncio.get_running_loop()
    tasks = []
    with concurrent.futures.ThreadPoolExecutor() as executor:
        for qube in appvms:
            if qube == default_dispvm and global_max is not None:
                maximum = global_max
                msg = f"global:{qube}:{maximum}"
            else:
                maximum = get_preload_max(qube)
                msg = f"{qube}:{maximum}"
            print(repr(msg))
            exec_args = qube.qubesd_call, qube.name, method, "preload-autostart"
            future = loop.run_in_executor(executor, *exec_args)
            tasks.append(future)
        await asyncio.gather(*tasks)


if __name__ == "__main__":
    asyncio.run(main())
