#!/bin/sh
imgdir=./debootstrap
default_sh_command=bash
default_chroot_command=/bin/bash
default_image_size=1500M # truncate(1) format

warn() { printf 'Warning: %s\n' "$*" >&2; }
die() { printf 'Error: %s\n' "$*" >&2; exit 1; }
[ -d "$imgdir" ] || die "directory does not exist: $imgdir"

usage()
{
    cat <<EOF >&2
Usage:

# creation & deletion
  $0 init <suite>
  $0 new <suite> <new-name>
  $0 clone <suite> <source-name> <new-name>
  $0 rm <suite> <source-name>

# modification
  $0 add      <suite> <name> <package> [packages...]
  $0 chroot   <suite> <name> [command] [args...]
  $0 sh       <suite> <name> [command] [args...]
  $0 with_mp  <suite> <name> <command> [args...]
  $0 with_img <suite> <name> <command> [args...]

# inspection
  $0 list
  $0 show unpackaged <suite> <name>
  $0 show unpackaged-du <suite> <name> ["filtered"]

EOF
    printf "Currently running suite:\n  %s\n" "$(lsb_release -cs)"
    list
    exit 1
}

list()
{
    local f suite name header_printed=
    for f in $imgdir/*.btrfs; do
        [ -e "$f" ] || continue
        f=${f##*/}
        f=${f%.btrfs}
        case "$f" in *.*) continue ;; esac
        suite=${f%%-*}
        [ "$header_printed" ] || echo 'Existing initialized suites:'
        printf '  %s\n' "$suite"
        header_printed=y
    done
    header_printed=
    for f in $imgdir/*.*.btrfs; do
        [ -e "$f" ] || continue
        f=${f##*/}
        f=${f%.btrfs}
        suite=${f%%-*}
        name=${f#*.}
        [ "$header_printed" ] || echo 'Existing images:'
        printf '  %s %s\n' "$suite" "$name"
        header_printed=y
    done
}

validate_suite()
{
    case "$1" in
        jessie|stretch|sid) return 0 ;;
        *) return 1 ;;
    esac
}

i_am_root() { [ "$(id -u)" = 0 ]; }

btrfs_show_subvolume_id()
{
    local result path="$1"
    result=$(btrfs subvolume show "$path" | sed -n -e 's/^[ \t]*Subvolume ID:[ \t]*//p; s/.*is toplevel subvolume/5/p')
    if [ "$result" ]
    then printf '%s\n' "$result"
    else false
    fi
}

mkfs_btrfs()
{
    local target="$1"
    mountpoint="$2" # can't be local because of trap

    mkfs.btrfs "$target"                                    || die "mkfs.btrfs failed"
    mount -t btrfs "$target" "$mountpoint"                  || die "mount failed"
    trap 'umount "$mountpoint"' EXIT

    btrfs subvolume create "$mountpoint"/root               || die "command 'btrfs subvolume create' failed"
    mkdir -p "$mountpoint"/root/var/cache/apt
    btrfs subvolume create "$mountpoint"/root/var/cache/apt/archives || die "command 'btrfs subvolume create' failed"
    subvol_id=$(btrfs_show_subvolume_id "$mountpoint"/root)          || die "could not find btrfs subvolume name"
    btrfs subvolume set-default "$subvol_id" "$mountpoint"           || die "command 'btrfs subvolume set-default' failed"

    trap - EXIT
    umount "$mountpoint"
}

suite_to_basename()
{
    suite=$1
    variant=minbase
    arch=$(dpkg-architecture -q DEB_HOST_ARCH) || die "command 'dpkg-architecture' failed"

    printf '%s\n' "$suite-$variant-$arch.btrfs"
}

suite_name_to_imagename()
{
    suite=$1
    name=$2
    variant=minbase
    arch=$(dpkg-architecture -q DEB_HOST_ARCH) || die "command 'dpkg-architecture' failed"

    printf '%s/%s-%s-%s.%s.btrfs\n' "$imgdir" "$suite" "$variant" "$arch" "$name"
}

show_unpackaged_helper()
{
    suite=$1
    name=$2
    imagename=$(suite_name_to_imagename "$suite" "$name")

    sh_image "$suite" "$name" find -type f | sed 's?^\./??' | sort > ${imagename}.find.txt
    sh_image "$suite" "$name" sh -c 'sort -u var/lib/dpkg/info/*.list' | sed -e 's?^/??' -e 's?^\.$/??' > ${imagename}.dpkg-list.txt
    comm -23 ${imagename}.find.txt ${imagename}.dpkg-list.txt > ${imagename}.unpackaged.txt
}

show_unpackaged()
{
    show_unpackaged_helper "$@"
    cat "${imagename}.unpackaged.txt"
}

show_unpackaged_du()
{
    show_unpackaged_helper "$@"
    if [ "$3" = filtered ]; then
        filter='var/lib/dpkg/|var/lib/apt/|var/log/|var/cache/apt/archives'
        cat "${imagename}.unpackaged.txt" | egrep -v "^($filter)" | sh_image "$1" "$2" xargs du -csh
    else
        cat "${imagename}.unpackaged.txt" | sh_image "$1" "$2" xargs du -csh
    fi
}

with_img()
{
    local suite="$1" name="$2" command="$3"
    shift 3
    [ "$suite" -a "$name" ] || usage
    imagename=$(suite_name_to_imagename "$suite" "$name")
    [ -e "$imagename" ] || die "no such file: $imagename"

    "$command" "$@" "$imagename"
}

run_command_with_mountpoint()
{
    local suite="$1" name="$2" command="$3"
    shift 3
    [ "$suite" -a "$name" ] || usage
    imagename=$(suite_name_to_imagename "$suite" "$name")
    [ -e "$imagename" ] || die "no such file: $imagename"
    [ -d "$imagename".mnt ] || mkdir "$imagename".mnt || die "mkdir"
    mountpoint -q "$imagename".mnt || mount "$imagename" "$imagename".mnt || die "mount"

    "$command" "$imagename".mnt "$@"

    r=$?
    umount "$imagename".mnt
    rmdir "$imagename".mnt
    return $r
}

rotate_args_left()
{
    local fst="$1" snd="$2"
    shift 2
    "$snd" "$@" "$fst"
}

add()
{
    export SKIP_UPDATE=y
    suite=$1
    name=$2
    shift 2
    running_suite=$(lsb_release -cs) || die 'lsb_release failed'
    if [ "$suite" = "$running_suite" -a -x ./src/selfstrap ]; then
        with_mp "$suite" "$name" ./src/selfstrap "$@" -t
    else
        exit 1
    fi
}

with_mp()
{
    suite=$1
    name=$2
    shift 2
    if [ $# = 0 ]; then
        cmd=echo
    else
        cmd=$1
        shift
    fi
    run_command_with_mountpoint "$suite" "$name" rotate_args_left "$cmd" "$@"
}



cd_then_run()
{
    (cd "$1"; shift; "$@")
}

sh_image()
{
    suite=$1
    name=$2
    shift 2
    if [ $# = 0 ]; then
        set -- $default_sh_command
    fi
    run_command_with_mountpoint "$suite" "$name" cd_then_run "$@"
}

chroot_image()
{
    suite=$1
    name=$2
    shift 2
    [ $# -gt 0 ] || set -- $default_chroot_command
    cgdir=/sys/fs/cgroup/pids/vm-$suite-$name
    mkdir "$cgdir"
    set -- \
        cgexec -g pids:vm-$suite-$name \
        unshare -f -m -p \
                chroot . \
                /bin/sh -c 'mount -t proc proc /proc; mount -t devpts devpts /dev/pts; exec "$@"' sh \
                "$@"
    sh_image "$suite" "$name" "$@"

    kill_cgroup "$cgdir"
}

kill_cgroup()
{
    (
        cgdir="$1"

        # exec >/dev/null 2>&1

        kill=
        sleep=0.2
        for n in $(seq 1 20); do

            [ ! -d "$cgdir" ] || rmdir "$cgdir" && return

            pids=$(cat "$cgdir"/tasks)
            [ -z "$pids" ] || continue

            kill $kill $pids
            sleep $sleep
            [ "$n" -eq 5 ] && sleep=1
            [ "$n" -eq 7 ] && sleep=3
            [ "$n" -eq 10 ] && kill=-KILL
        done
        return 1
    )
}

clone()
{
    suite=$1
    source_name=$2
    name=$3
    [ "$suite" -a "$name" -a "$source_name" ] || usage
    image_basename=$imgdir/$(suite_to_basename "$suite") || die 'suite_to_basename'

    source=${image_basename%.btrfs}.$source_name.btrfs
    dest=${image_basename%.btrfs}.$name.btrfs

    [ -e "$source" ] || die "source image not found -- run '$0 new $suite $source_name' to create it"
    [ ! -e "$dest" ] || die "file exists: $dest"
    cp --reflink=always "$source" "$dest"
}

new()
{
    suite=$1
    name=$2
    [ "$suite" -a "$name" ] || usage
    image_basename=$imgdir/$(suite_to_basename "$suite") || die 'suite_to_basename'

    source=$image_basename
    dest=${source%.btrfs}.$name.btrfs

    [ -e "$source" ] || die "source image not found -- run '$0 init $suite' to create it"
    [ ! -e "$dest" ] || die "file exists: $dest"
    cp --reflink=always "$source" "$dest"
}

rm_image()
{
    suite=$1
    name=$2
    [ "$suite" -a "$name" ] || usage
    image_basename=$imgdir/$(suite_to_basename "$suite") || die 'suite_to_basename'
    source=$image_basename
    dest=${source%.btrfs}.$name.btrfs

    rm -i "$dest"
}

init()
{
    suite=$1
    variant=minbase
    arch=$(dpkg-architecture -q DEB_HOST_ARCH) || die "command 'dpkg-architecture' failed"
    running_suite=$(lsb_release -cs) || die 'lsb_release failed'

    size=$default_image_size

    i_am_root || die 'you are not root'

    validate_suite "$suite" || die "invalid suite: '$suite'"

    target=$imgdir/$suite-$variant-$arch.btrfs
    [ -e "$target" ] && return
    [ -e "$target".tmp ] && die "refusing to overwrite existing temporary file: $target.tmp"

    [ -d "$target".mnt ] || mkdir "$target".mnt || die "could not create directory $target.mnt"

    truncate -s "$size" "$target".tmp          || die "truncate failed"
    mkfs_btrfs "$target".tmp "$target".mnt     || die "btrfs filesystem creation failed"
    mount -t btrfs "$target".tmp "$target".mnt || die "mount failed"
    trap 'umount "$target.mnt"' EXIT

    if [ "$suite" = "$running_suite" -a -x ./src/selfstrap ]; then
        # TODO:
        # 1. Use a --chroot arg that will set up cgroups
        # 2. Use --packages arg to implement variant
        # 3. Allow selfstrap to run on already-initialized image
        ./src/selfstrap -t "$target.mnt" || die "selfstrap failed"
    else
        debootstrap_efficiently "$arch" "$variant" "$suite" \
                                "$target".mnt "${target%.btrfs}".debs.txt || die "debootstrap failed"
    fi

    if umount "$target".mnt
    then trap - EXIT
    else warn "umount failed"
    fi

    mv "$target".tmp "$target"
}

debootstrap_efficiently()
{
    local arch="$1" variant="$2" suite="$3" target="$4" savedebs="$5"
    set -- --arch "$arch" --variant "$variant" "$suite" "$target"
    debs=$(debootstrap --print-debs --keep-debootstrap-dir "$@" | tee "$savedebs") || die "debootstrap failed"
    for deb in $debs; do
        cp /var/cache/apt/archives/${deb}_* $target/var/cache/apt/archives/
    done

    debootstrap "$@" || die "debootstrap failed"
}

show()
{
    case "$1" in
        unpackaged) shift; show_unpackaged "$@" ;;
        unpackaged-du) shift; show_unpackaged_du "$@" ;;
        *) usage ;;
    esac
}

case "$1" in
    init|new|clone|list|show|add|with_mp|with_img) cmd=$1; shift; $cmd "$@" ;;
    chroot|sh|rm) cmd=${1}_image; shift; $cmd "$@" ;;
    *) usage ;;
esac