{ pkgs, ... }:
pkgs.writers.writePython3 "screenshot" {  flakeIgnore = [ "E501" ];} ''
import subprocess
import re

options = """\
a selected area (copy)
current window (copy)
full screen (copy)
a selected area
current window
full screen
viewer
"""

cmd = '${pkgs.rofi-wayland}/bin/rofi -dmenu -location 2 -l 7 -i -p "Screenshot which area?"'
result = subprocess.run(cmd,
                        input=options,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        shell=True,
                        text=True)

if result.returncode != 0:
    exit(1)

choice = result.stdout.strip()
cmd = "${pkgs.pkgs.sway-contrib.grimshot}/bin/grimshot "
if choice == "viewer":
    filename = subprocess.getoutput("mktemp")
    cmd += f"save area {filename}"
else:
    cmd += "copy " if re.match(".*(copy)", choice) else "save "

for k, v in {"a selected area.*": "area ",
             "current window.*": "active ",
             "current screen.*": "output ",
             "all screens.*": "screen "}.items():
    if re.match(k, choice):
        cmd += v

subprocess.call(cmd,
                shell=True)

if choice == "viewer":
    subprocess.run(f"nsxiv -b {filename}", shell=True)
''