#!/usr/bin/python

import asyncio
import json
import os
import subprocess
import sys


def get_outputs():
    """
       Returns the list of active sway outputs
    """
    cmd = [
        'swaymsg',
        '-t',
        'get_outputs',
    ]
    outputs = json.loads(subprocess.check_output(cmd))
    return [ out['name'] for out in outputs if out['active'] ]


async def grab_and_blur(display, path):
    """
       Will screenshot+blur `display` at `path/display.png`
    """
    output = os.path.join(path, '{}.png'.format(display))
    screengrab = subprocess.Popen(['grim', '-o', display, '-'], stdout=subprocess.PIPE)
    blur = await asyncio.subprocess.create_subprocess_exec('ffmpeg',
            '-loglevel', 'quiet','-y',
            '-i', '-', '-vframes', '1', '-filter_complex', 'boxblur=10:2',
            output,
            stdin=screengrab.stdout)
    await blur.wait()
    # screengrab should already be dead
    screengrab.wait()

    return output


async def main():
    # Would be tmpdir but can't cleanup because of `execv`
    dst_dir = '/tmp'
    lock_argv = ['swaylock']
    outputs = get_outputs()
    grab_routines = []

    for display in outputs:
        grab_routines.append(grab_and_blur(display, dst_dir))

    images = await asyncio.gather(*grab_routines)

    for i, display in enumerate(outputs):
        lock_argv.extend([
            '-i',
            f'{display}:{images[i]}',
        ])

    lock_argv.extend(sys.argv[1:])
    return lock_argv

if __name__ == '__main__':
    lock_cmd = asyncio.run(main())
    os.execvp(lock_cmd[0], lock_cmd)
