#!/usr/bin/env python3
from functools import partial
from itertools import chain
import argparse
import os
import shutil
import subprocess


def flat_map(f, iterable):
    return list(chain.from_iterable(map(f, iterable)))


def add_switch(name: str, default=False):
    if default:
        parser.add_argument(f"--no-{name}", dest=name, action="store_false")
    else:
        parser.add_argument(
            f"--{name}", dest=name, action="store_true", default=default
        )


def tmp_file(name: str):
    tmpdir = f"/tmp/bwrap-helper-{os.getpid()}"
    os.makedirs(tmpdir, exist_ok=True)
    if name is None:
        return tmpdir
    else:
        return f"{tmpdir}/{name}"


def generate_tmp_file(name: str, content: str):
    file_path = tmp_file(name)
    with open(file_path, "w") as f:
        f.write(content)
    return file_path


def bind(src: str, dest=None, type=None, required=True):
    arg = "--"
    if type is not None:
        arg += type + "-"
    arg += "bind"
    if not required:
        arg += "-try"

    if dest is None:
        dest = src

    return [arg, src, dest]


def setenv(name: str, value: str):
    return ["--setenv", name, value]


def parse_passthrough_arg(name: str, nargs: int):
    parser.add_argument(f"--{name}", action="append", nargs=nargs)


def assemble_passthrough_arg(name: str):
    for value in getattr(args, name.replace("-", "_")) or []:
        assembled_args.extend([f"--{name}", *value])


dev_bind = partial(bind, type="dev")
ro_bind = partial(bind, type="ro")
ro_bind_try = partial(bind, type="ro", required=False)

username = os.getenv("USER")
uid = os.getuid()
gid = os.getgid()
home = os.getenv("HOME")

argument_groups = {
    "base": (
        True,
        [
            "--tmpfs",
            "/tmp",
            "--proc",
            "/proc",
            "--dev",
            "/dev",
            "--dir",
            home,
            "--dir",
            f"/run/user/{uid}",
            *ro_bind("/etc/localtime"),
            "--unshare-all",
            "--die-with-parent",
        ],
    ),
    "nix-store": (
        True,
        [
            *flat_map(
                ro_bind,
                [
                    "/nix/store",
                    "/etc/static",
                ],
            ),
        ],
    ),
    "path": (
        True,
        [
            *flat_map(
                ro_bind,
                [
                    "/run/current-system/sw",  # not exclusive to path, but also libraries etc.
                    f"/etc/profiles/per-user/{username}/bin",
                ],
            ),
        ],
    ),
    "gui": (
        False,
        [
            *dev_bind("/dev/dri"),
            *flat_map(
                ro_bind,
                [
                    "/sys/dev/char",
                    "/sys/devices/pci0000:00",
                    f"/run/user/{uid}/{os.getenv('WAYLAND_DISPLAY')}",
                    "/run/opengl-driver",
                    "/etc/fonts",
                ],
            ),
            *ro_bind_try("/run/opengl-driver-32"),
        ],
    ),
    "x11": (
        False,
        [
            *ro_bind("/tmp/.X11-unix"),
        ],
    ),
    "audio": (
        False,
        [
            *ro_bind(f"/run/user/{uid}/pulse"),
            # should in theory autodetect, but sometimes it does not work
            *setenv("PULSE_SERVER", f"/run/user/{uid}/pulse/native"),
            # some programs need the cookie
            *ro_bind(f"{home}/.config/pulse/cookie"),
            *setenv("PULSE_COOKIE", f"{home}/.config/pulse/cookie"),
            # ALSA compat
            *ro_bind("/etc/asound.conf", required=False),
            *ro_bind("/etc/alsa/conf.d", required=False),
            # pipewire
            *ro_bind(f"/run/user/{uid}/pipewire-0", required=False),
        ],
    ),
    "passwd": (
        False,
        [
            *ro_bind(
                generate_tmp_file(
                    "passwd",
                    f"{username}:x:{uid}:{gid}::{home}:/run/current-system/sw/bin/bash\n",
                ),
                "/etc/passwd",
            )
        ],
    ),
    "network": (
        False,
        [
            "--share-net",
            *flat_map(
                ro_bind,
                [
                    "/etc/resolv.conf",
                    "/etc/ssl/certs",
                ],
            ),
        ],
    ),
    "dbus": (
        False,
        [
            *ro_bind("/run/dbus/system_bus_socket"),
            *ro_bind(generate_tmp_file("machine-id", "0" * 32), "/etc/machine-id"),
        ],
    ),
    "new-session": (
        True,
        [
            "--new-session",
        ],
    ),
    "pwd": (
        False,
        [
            *ro_bind(os.getcwd()),
            "--chdir",
            os.getcwd(),
        ],
    ),
    "pwd-rw": (
        False,
        [
            *bind(os.getcwd()),
            "--chdir",
            os.getcwd(),
        ],
    ),
}

passthrough_args = [
    ("bind", 2),
    ("dev-bind", 2),
    ("dev-bind-try", 2),
    ("ro-bind", 2),
    ("symlink", 2),
]

for (_, arguments) in argument_groups.values():
    for argument in arguments:
        assert type(argument) == str

parser = argparse.ArgumentParser()
parser.add_argument("--show-cmdline", action="store_true")
for name, (default, _) in argument_groups.items():
    add_switch(name, default)
parser.add_argument("program")
parser.add_argument("args", nargs="*")

for (arg, nargs) in passthrough_args:
    parse_passthrough_arg(arg, nargs)

args = parser.parse_args()

assembled_args = ["bwrap"]

for name, (_, arguments) in argument_groups.items():
    if getattr(args, name):
        assembled_args.extend(arguments)

for (arg, _) in passthrough_args:
    assemble_passthrough_arg(arg)

if args.show_cmdline:
    for idx, assembled_arg in enumerate(assembled_args):
        if idx == 0:
            print(assembled_arg, end="")
            continue
        if assembled_arg.startswith("--"):
            print("\n    ", end="")
        else:
            print(end=" ")
        print(assembled_arg, end="")
    print()

assembled_args.append(args.program)
assembled_args.extend(args.args)

try:
    subprocess.run(assembled_args)
finally:
    shutil.rmtree(tmp_file(None))