#!/bin/bash
Usage() {
   cat <<EOF
${0##*/} out-directory release
example:
 ${0##*/} trusty_dir trusty daily
EOF
}
TEMP_D=""

cleanup() {
   [ ! -d "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}
VERBOSITY=0
inargs() {
  local n=$1 x="";
  shift;
  for x in "$@"; do [ "$n" = "$x" ] && return 0; done;
  return 1
}
debug() { [ $1 -lt $VERBOSITY ] || return 0; shift; echo "$@" 1>&2; }
error() { echo "$@" 1>&2; }
fail() { [ $# -eq 0 ] || error "$@"; exit 1; }

VERBOSITY=0
TEMP_D=$(mktemp -d ${TMPDIR:-/tmp}/${0##*/}.XXXXXX) || exit 1
trap cleanup EXIT

arches=( ppc64el i386 amd64 )
releases=( $(ubuntu-distro-info --all) )
mbase="http://maas.ubuntu.com/images/ephemeral-v2"
def_stream="${mbase}/daily"
def_rel="trusty"
vername=""

[ "$1" = "-h" -o "$1" = "--help" ] && { Usage; exit 0; }
[ $# -lt 2 ] && { Usage 1>&2; exit 1; }
out_d="$1"
shift

myarch=$(uname -m)
case "$myarch" in
   ppc64*) def_arch="ppc64el";;
   i?86) def_arch="i386";;
   x86_64) def_arch="amd64";;
   arm*) def_arch="armel";;
esac

pt=( )
for x in "$@"; do
   inargs "$x" "${arches[@]}" && pt[${#pt[@]}]=arch=$x && arch="$x" && continue
   inargs "$x" "${releases[@]}" && pt[${#pt[@]}]="release=$x" && release=$x && continue
   case "$x" in
      daily) stream="${mbase}/daily";;
      released|release|releases) stream="${mbase}/releases";;
      http://*) stream=$x;;
      hwe-*) subarch="$x"; pt[${#pt[@]}]="subarch=$x";;
      subarch=*) pt[${#pt[@]}]="$x"; subarch=${x#*=};;
      *=*) pt[${#pt[@]}]="$x";;
      *~*) pt[${#pt[@]}]="$x";;
      [0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9]|[0-9][0-9][0-9][0-9][0-9][0-9][0-9][0-9].[0-9])
         pt[${#pt[@]}]="version_name=$x"
         vername="$x";;
   esac
done

[ -n "$stream" ] || stream="$def_stream"
[ -n "$arch" ] || pt[${#pt[@]}]="arch=$def_arch"
if [ -z "$subarch" ]; then
  t=${release#?}
  first_letter=${release%${t}}
  subarch="hwe-${first_letter}"
  pt[${#pt[@]}]="subarch=$subarch"
  error "selected subarch=${subarch} for $release"
fi

OIFS="$IFS"
mkdir -p "$out_d"
case "$stream" in
   *.json) :;;
   *) keyring="/usr/share/keyrings/ubuntu-cloudimage-keyring.gpg";;
esac

needs=""
[ -n "$keyring" -a ! -f "$keyring" ] &&
   needs="${needs} ubuntu-cloudimage-keyring"

for pair in sstream-query:simplestreams tgt-admin:tgt; do
   cmd=${pair%:*}
   pkg=${pair#*:}
   command -v "$cmd" >/dev/null 2>&1 || needs="$needs $pkg"
done

if [ -n "${needs# }" ]; then
   error "missing dependencies"
   fail "sudo apt-get install -qy ${needs# }"
fi

qcmd=( sstream-query ${keyring:+"--keyring=$keyring"} ${vername:---max=1}
       --output-format="%(sha256)s %(subarch)s %(ftype)s %(item_url)s" "${stream}" "${pt[@]}" )
echo "${qcmd[@]}"
"${qcmd[@]}" > "${TEMP_D}/qout"
roots=$(awk '$3 == "root-image.gz" { print $2 }' "${TEMP_D}/qout" | sort -u | wc -l)
[ "$roots" = "1" ] || {
   error "query resulted in '$roots' root images. expect 1 and only 1"
   cat "${TEMP_D}/qout"
   fail
}
image="$out_d/root-image"
set -f
while read line; do
    set -- $line
    sha256="$1"; subarch="$2"; ftype="$3"; url="$4"
    [ "$ftype" = "di-kernel" -o "$ftype" = "di-initrd" ] && continue
    if [ "$ftype" = "root-image.gz" ]; then
       fname="$ftype"
       [ -f "$image" ] && continue
    else
       fname="$subarch/$ftype"
    fi
    out_f="${out_d}/$fname"
    [ -f "$out_f" ] && continue
    [ -d "${out_f%/*}" ] || mkdir -p "${out_f%/*}"
    wget -c "$url" -O "$out_f.tmp" && mv "$out_f.tmp" "$out_f" || { rm -f "$out_f.tmp"; exit 3; }
    echo "$out_d/$fname" "$url" >> "$out_d/contents"
done < "${TEMP_D}/qout"

if [ ! -f "$image" ]; then
   debug 1 "decompressing $image.gz"
   zcat "$image.gz" > "$image.tmp" &&
     mv "$image.tmp" "$image" ||
     { rm -f "$image.tmp"; exit 1; }
fi

# vi: ts=4 expandtab
