#!/usr/bin/env bash
SCRIPT_PATH=$(dirname $(realpath -s $0))
MUVM_PATH=$(dirname $(which muvm))
PASST_PATH=$(dirname $(which passt))
HOST_OPENGL_DRIVER=/run/opengl-driver
MICROVM_CLOSURE=
MICROVM_COMMAND=()
MICROVM_UID=1337
MICROVM_GID=1337
BWRAP_ARGS=()
MUVM_ARGS=()
GPU=1
WAYLAND=1
PIPEWIRE=1
X11=0
export TMP=/tmp TMPDIR=/tmp TEMP=/tmp TEMPDIR=/tmp LC_ALL=C

while [ "$#" -gt 0 ]; do
  case "$1" in
    -u|--uid) MICROVM_UID="$2"; shift 2;;
    -g|--gid) MICROVM_GID="$2"; shift 2;;
    --no-gpu) GPU=0; shift 1;;
    --no-wayland) WAYLAND=0; shift 1;;
    --no-pipewire) PIPEWIRE=0; shift 1;;
    --x11) X11=1; shift 1;;
    --bind) BWRAP_ARGS+=("--bind" "$2" "$3"); shift 3;;
    --ro-bind) BWRAP_ARGS+=("--ro-bind" "$2" "$3"); shift 3;;
    --expose) BWRAP_ARGS+=("--bind" "$2" "$2"); shift 2;;
    --ro-expose) BWRAP_ARGS+=("--ro-bind" "$2" "$2"); shift 2;;
    --host-opengl-driver) HOST_OPENGL_DRIVER="$2"; shift 2;;
    --munix-bin-dir) SCRIPT_PATH="$2"; shift 2;;
    --muvm-bin-dir) MUVM_PATH="$2"; shift 2;;
    --passt-bin-dir) PASST_PATH="$2"; shift 2;;
    -*) echo "munix: unknown option: $1" >&2; exit 1;;
    *)
      if [ "$MICROVM_CLOSURE" = "" ]; then
        MICROVM_CLOSURE="$1"
      else
        MICROVM_COMMAND+=("$1")
      fi
      shift 1;;
  esac
done

if [ "$MICROVM_CLOSURE" = "" ]; then
  echo "munix: provide a system closure path as a positional argument" >&2
  exit 1
fi

if [ ${#MICROVM_COMMAND[@]} -eq 0 ]; then
  MICROVM_COMMAND=("bash")
fi

if [ "$GPU" -eq 1 ]; then
  BWRAP_ARGS+=(
    "--dev-bind" "/dev/dri" "/dev/dri"
    "--ro-bind" "$HOST_OPENGL_DRIVER" "/run/opengl-driver"
  )
  GPU_MODE=venus
  kernel_ver="$(uname -r)"
  kernel_ver_arr=(${kernel_ver//./ })
  kernel_major="${kernel_ver_arr:-0}"
  kernel_ver_arr=("${kernel_ver_arr[@]:1}")
  kernel_minor="${kernel_ver_arr:-0}"
  if [[ "$kernel_major" -gt 6 || ("$kernel_major" -eq 6 && "$kernel_minor" -gt 12) ]]; then
    for card in /dev/dri/card*; do
      driver_link="/sys/class/drm/${card##*/}/device/driver"
      if [ -L "$driver_link" ]; then
        driver_mod="$(readlink "$driver_link")"
        driver_name="${driver_mod##*/}"
        case "$driver_name" in
          amdgpu|msm) # TODO: i915
            echo "munix: ${card##*/} gpu driver is '$driver_name', using vdrm" >&2;
            GPU_MODE=drm
            break;;
          *) echo "munix: ${card##*/} gpu driver is '$driver_name', using venus unless more gpus are found" >&2;;
        esac
      else
        echo "munix: ${card##*/} has no gpu driver" >&2;
      fi
    done
  else
    echo "munix: kernel version '$kernel_ver' is older than 6.13, not using gpu due to missing support" >&2;
    GPU_MODE=software
  fi
  MUVM_ARGS+=("--gpu-mode=$GPU_MODE")
  if [ "$GPU_MODE" = "venus" ]; then
    MUVM_ARGS+=("-e" "MESA_LOADER_DRIVER_OVERRIDE=zink")
  fi
else
  BWRAP_ARGS+=("--dir" "/dev/dri")
  MUVM_ARGS+=("--gpu-mode=software")
fi

if [ "$WAYLAND" -eq 1 ]; then
  if [ "$XDG_RUNTIME_DIR" = "" ]; then
    echo "munix: wayland requested, but no XDG_RUNTIME_DIR set" >&2
    exit 1
  fi
  BWRAP_ARGS+=(
    "--bind" "$XDG_RUNTIME_DIR/$WAYLAND_DISPLAY" "$XDG_RUNTIME_DIR/$WAYLAND_DISPLAY"
    "--setenv" "WAYLAND_DISPLAY" "$WAYLAND_DISPLAY"
  )
  MUVM_ARGS+=("-e" "WAYLAND_DISPLAY=wayland-1") # the proxy is managed by us, not muvm
fi

if [ "$PIPEWIRE" -eq 1 ]; then
  if [ "$PIPEWIRE_RUNTIME_DIR" = "" ]; then
    PIPEWIRE_RUNTIME_DIR="$XDG_RUNTIME_DIR"
  fi
  if [ "$PIPEWIRE_RUNTIME_DIR" = "" ]; then
    PIPEWIRE_RUNTIME_DIR="$USERPROFILE"
  fi
  if [ "$PIPEWIRE_RUNTIME_DIR" = "" ]; then
    echo "munix: pipewire requested, but no PIPEWIRE_RUNTIME_DIR/XDG_RUNTIME_DIR/USERPROFILE set" >&2
    exit 1
  fi
  if [ "$PIPEWIRE_REMOTE" = "" ]; then
    PIPEWIRE_REMOTE=pipewire-0
  fi
  BWRAP_ARGS+=("--bind" "$PIPEWIRE_RUNTIME_DIR/$PIPEWIRE_REMOTE" "$PIPEWIRE_RUNTIME_DIR/$PIPEWIRE_REMOTE")
fi

if [ "$X11" -eq 1 ]; then
  BWRAP_ARGS+=(
    "--bind" "/tmp/.X11-unix" "/tmp/.X11-unix"
    "--ro-bind" "$XAUTHORITY" "$XAUTHORITY"
  )
else
  unset DISPLAY XAUTHORITY
fi

exec bwrap --unshare-all --share-net \
  --uid $MICROVM_UID --gid $MICROVM_GID \
  --tmpfs / \
  --dir /run --dir /var --symlink /run /var/run --dir /tmp \
  --proc /proc --ro-bind /sys /sys \
  --dev /dev --dir /dev/input --dev-bind /dev/kvm /dev/kvm \
  --ro-bind "$MUVM_PATH" /run/munix/muvm \
  --ro-bind "$PASST_PATH" /run/munix/passt \
  --ro-bind "$MICROVM_CLOSURE/sw/bin/env" /usr/bin/env \
  --ro-bind "$SCRIPT_PATH/munix-init-root" /usr/bin/munix-init-root \
  --ro-bind "$SCRIPT_PATH/munix-init-user" /usr/bin/munix-init-user \
  --ro-bind /nix/store /nix/store \
  --ro-bind /run/systemd/resolve /run/systemd/resolve \
  --ro-bind /etc/resolv.conf /etc/resolv.conf \
  --file 11 /etc/passwd \
  --file 12 /etc/group \
  --dir "$XDG_RUNTIME_DIR" \
  --setenv PATH "/run/munix/muvm:/run/munix/passt:$MICROVM_CLOSURE/sw/bin" \
  "${BWRAP_ARGS[@]}" \
  muvm \
    -x /usr/bin/munix-init-root -X /usr/bin/munix-init-user --udevd-path="$MICROVM_CLOSURE/sw/bin/true" \
    "${MUVM_ARGS[@]}" \
    -e MICROVM_CLOSURE="$MICROVM_CLOSURE" -e MICROVM_UID="$MICROVM_UID" -e MICROVM_GID="$MICROVM_GID" \
    -i -t "${MICROVM_COMMAND[@]}" \
  11< <(cat <<EOF
munix:x:$MICROVM_UID:$MICROVM_GID:Hypervisor:/:/run/current-system/sw/bin/nologin
nobody:x:65534:65534:Unprivileged account:/var/empty:/run/current-system/sw/bin/nologin
EOF
) \
  12< <(cat <<EOF
munix:x:$MICROVM_GID:
nogroup:x:65534:
EOF
)