# vi: ts=4 expandtab syntax=sh

#imagesize=${IMAGE_SIZE:-$((2252*1024**2))}  # 2.2G (the current size we ship)
imagesize=${IMAGE_SIZE:-2361393152}  # 2.2G (the current size we ship)

release_ver() {
    # Return the release version number
    distro-info --series="$LB_DISTRIBUTION" -r | awk '{ print $1 }'
}

_snap_preseed() {
    # Download the snap/assertion and add to the preseed
    local CHROOT_ROOT=$1
    local SNAP=$2
    local SNAP_NAME=${SNAP%/*}
    local CHANNEL=${3:?Snap channel must be specified}

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"
    local seed_yaml="$seed_dir/seed.yaml"
    local assertions_dir="$seed_dir/assertions"

    # Download the snap & assertion
    local snap_download_failed=0
    chroot $CHROOT_ROOT sh -c "
        set -x;
        cd /var/lib/snapd/seed;
        SNAPPY_STORE_NO_CDN=1 snap download \
            --channel=$CHANNEL \"$SNAP_NAME\"" || snap_download_failed=1
    if [ $snap_download_failed = 1 ] ; then
        echo "If the channel ($CHANNEL) includes '*/ubuntu-##.##' track per "
        echo "Ubuntu policy (ex. stable/ubuntu-18.04) the publisher will need "
        echo "to temporarily create the channel/track to allow fallback during"
        echo "download (ex. stable/ubuntu-18.04 falls back to stable if the"
        echo "prior had been created in the past)."
        exit 1
    fi

    mv -v $seed_dir/*.assert $assertions_dir
    mv -v $seed_dir/*.snap $snaps_dir

    # Add the snap to the seed.yaml
    ! [ -e $seed_yaml ] && echo "snaps:" > $seed_yaml
    cat <<EOF >> $seed_yaml
  -
    name: ${SNAP_NAME}
    channel: ${CHANNEL}
EOF

    case ${SNAP} in */classic) echo "    classic: true" >> $seed_yaml;; esac

    echo -n "    file: " >> $seed_yaml
    (cd $snaps_dir; ls -1 ${SNAP_NAME}_*.snap) >> $seed_yaml
}

snap_prepare_assertions() {
    # Configure basic snapd assertions
    local CHROOT_ROOT=$1
    # A colon-separated string of brand:model to be used for the image's model
    # assertion
    local CUSTOM_BRAND_MODEL=$2

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"
    local assertions_dir="$seed_dir/assertions"
    local model_assertion="$assertions_dir/model"
    local account_key_assertion="$assertions_dir/account-key"
    local account_assertion="$assertions_dir/account"

    mkdir -p "$assertions_dir"
    mkdir -p "$snaps_dir"

    local brand="$(echo $CUSTOM_BRAND_MODEL | cut -d: -f 1)"
    local model="$(echo $CUSTOM_BRAND_MODEL | cut -d: -f 2)"

    if ! [ -e "$model_assertion" ] ; then
        snap known --remote model series=16 \
            model=$model brand-id=$brand \
            > "$model_assertion"
    fi

    if ! [ -e "$account_key_assertion" ] ; then
        local account_key=$(sed -n -e's/sign-key-sha3-384: //p' \
            < "$model_assertion")
        snap known --remote account-key \
            public-key-sha3-384="$account_key" \
            > "$account_key_assertion"
    fi


    if ! [ -e "$account_assertion" ] ; then
        local account=$(sed -n -e's/account-id: //p' < "$account_key_assertion")
        snap known --remote account account-id=$account \
            > "$account_assertion"
    fi
}

snap_prepare() {
    # Configure basic snapd assertions and pre-seeds the 'core' snap
    local CHROOT_ROOT=$1
    # Optional. If set, should be a colon-separated string of brand:model to be
    # used for the image's model assertion
    local CUSTOM_BRAND_MODEL=${2:-generic:generic-classic}

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"

    snap_prepare_assertions "$CHROOT_ROOT" "$CUSTOM_BRAND_MODEL"

    # Download the core snap
    if ! [ -f $snaps_dir/core_[0-9]*.snap ] ; then
        _snap_preseed $CHROOT_ROOT core stable
    fi
}

snap_preseed() {
    # Preseed a snap in the image
    local CHROOT_ROOT=$1
    local SNAP=$2
    # Per Ubuntu policy, all seeded snaps (with the exception of the core
    # snap) must pull from stable/ubuntu-$(release_ver) as their channel.
    local CHANNEL=${3:-"stable/ubuntu-$(release_ver)"}

    snap_prepare $CHROOT_ROOT
    _snap_preseed $CHROOT_ROOT $SNAP $CHANNEL
}