#!/bin/bash
set -e
cd "$(dirname "$0")"

#
# This script builds the neko base image and all the applications
#

function log() {
  echo "$(date +'%Y-%m-%d %H:%M:%S') - [NEKO] - $1" > /dev/stderr
}

function help() {
  echo "Usage: $0"
  echo "  -p, --platform : The platform (default: linux/amd64)"
  echo "  -i, --image    : The image name prefix (default: ghcr.io/m1k1o/neko)"
  echo "  -t, --tag      : The image tag (default: latest)"
  echo "  -f, --flavor   : The flavor, if not specified, builds without flavor"
  echo "  -b, --base     : The base image name (default: ghcr.io/m1k1o/neko/[<flavor>-]base:<tag>)"
  echo "  -a, --app      : The app to build, if not specified, builds the base image"
  echo "  -y, --yes      : Skip confirmation prompts"
  echo "      --no-cache : Build without docker cache"
  echo "      --push     : Push the image to the registry after building"
  echo "  -h, --help     : Show this help message"
}

FULL_IMAGE=""
while [[ "$#" -gt 0 ]]; do
  case $1 in
    --platform|-p) PLATFORM="$2"; shift ;;
    --image|-i) IMAGE="$2"; shift ;;
    --tag|-t)
      # comma separated list of TAGS
      IFS=',' read -r -a TAGS <<< "$2"
      # default tag is the first one
      TAG="${TAGS[0]}"
      shift ;;
    --flavor|-f) FLAVOR="$2"; shift ;;
    --base|-b) BASE_IMAGE="$2"; shift ;;
    --app|-a) APPLICATION="$2"; shift ;;
    --yes|-y) YES=1 ;;
    --no-cache) NO_CACHE="--no-cache" log "Building without cache" ;;
    --push) PUSH=1 ;;
    --help|-h) help; exit 0 ;;
    -*) log "Unknown parameter passed: $1"; help; exit 1 ;;
    *) 
      if [ -z "$FULL_IMAGE" ]; then
        FULL_IMAGE="$1"

        # extracts image, flavor, app and tag from the full image name
        # example:
        #  ghcr.io/m1k1o/neko/nvidia-firefox:latest
        # will be split into:
        #   IMAGE=ghcr.io/m1k1o/neko
        #   FLAVOR=nvidia
        #   APPLICATION=firefox
        #   TAG=latest
        
        # remove the tag from the image name
        if [[ "$FULL_IMAGE" == *":"* ]]; then
          # removes everything before the last :
          TAG="${FULL_IMAGE##*:}" # will be latest
          # removes everything after the last :
          FULL_IMAGE="${FULL_IMAGE%:*}" # will be ghcr.io/m1k1o/neko/nvidia-firefox
        fi
        # extract the image name and save the rest to IMAGE
        if [[ "$FULL_IMAGE" == *"/"* ]]; then
          # removes everything after the last /
          IMAGE="${FULL_IMAGE%/*}" # will be ghcr.io/m1k1o/neko
          # removes everything before the last /
          FULL_IMAGE="${FULL_IMAGE##*/}" # will be nvidia-firefox
        fi
        # extract the flavor and application name
        if [[ "$FULL_IMAGE" == *"-"* ]]; then
          # removes everything after the last -
          FLAVOR="${FULL_IMAGE%-*}" # will be nvidia
          # removes everything before the last -
          APPLICATION="${FULL_IMAGE#*-}" # will be firefox
        else
          # no flavor specified so use the full image name as application name
          APPLICATION="$FULL_IMAGE" # will be firefox
        fi
        # if application name is base, set it to empty
        if [ "$APPLICATION" == "base" ]; then
          APPLICATION=""
        fi
      else
        log "Unknown positional argument: $1"
        help
        exit 1
      fi
    ;;
  esac
  shift
done

function prompt() {
  if [ ! -z "$YES" ]; then
    return 0
  fi

  local OK=""
  while [ -z "$OK" ]; do
    read -p "$1 (yes/no) " REPLY
    case "$REPLY" in
      yes|YES|y|Y)
        OK=1
        ;;
      no|NO|n|N)
        log "Aborting build."
        exit 1
        ;;
      *)
        log "Please answer 'yes' or 'no'."
        ;;
    esac
  done
}

function build_image() {
  docker build --platform $PLATFORM $NO_CACHE $@
}

function tag_image() {
  local IMAGE="$1"
  local IMAGE_NO_TAG="${IMAGE%:*}"
  local IMAGE_TAG="${IMAGE##*:}"

  for TAG in "${TAGS[@]}"; do
    # skip if the tag is the same as the image tag
    if [ "$TAG" == "$IMAGE_TAG" ]; then
      continue
    fi

    log "Tagging $IMAGE with tag $TAG"
    docker tag $IMAGE $IMAGE_NO_TAG:$TAG
  done
}

function push_image() {
  if [ -z "$PUSH" ]; then
    return 0
  fi

  local IMAGE="$1"
  local IMAGE_NO_TAG="${IMAGE%:*}"

  for TAG in "${TAGS[@]}"; do
    log "Pushing $IMAGE_NO_TAG:$TAG to registry"
    docker push $IMAGE_NO_TAG:$TAG
  done
}

# --------------------------------------------------------------------

if [ -z "$PLATFORM" ]; then
  PLATFORM="linux/amd64"
fi
log "Using platform: $PLATFORM"

if [ -z "$IMAGE" ]; then
  IMAGE="ghcr.io/m1k1o/neko"
fi
log "Using image: $IMAGE"

if [ -z "$TAG" ]; then
  TAG="latest"
  TAGS=($TAG)
fi

for TAG in "${TAGS[@]}"; do
  log "Using tag: $TAG"
done

if [ -z "$FLAVOR" ]; then
  log "No flavor specified, building without flavor"
else
  log "Using flavor: $FLAVOR"
fi

if [ -z "$BASE_IMAGE" ]; then
  if [ -z "$FLAVOR" ]; then
    BASE_IMAGE="$IMAGE/base:$TAG"
  else
    BASE_IMAGE="$IMAGE/$FLAVOR-base:$TAG"
  fi
fi

# --------------------------------------------------------------------

if [ ! -z "$APPLICATION" ]; then
  log "Building application: $APPLICATION"
  log "Using base image: $BASE_IMAGE"

  # check if application directory exists
  APPLICATION_DIR="apps/$APPLICATION"
  if [ ! -d "$APPLICATION_DIR" ]; then
    log "Application directory $APPLICATION_DIR does not exist."
    exit 1
  fi

  # flavor is specified, append it to the image name and Dockerfile
  APPLICATION_IMAGE="$IMAGE/$APPLICATION:$TAG"
  APPLICATION_DOCKERFILE="apps/$APPLICATION/Dockerfile"
  if [ ! -z "$FLAVOR" ]; then
    APPLICATION_IMAGE="$IMAGE/$FLAVOR-$APPLICATION"
    # if application flavor is specified and Dockerfile exists, use it
    if [ -f "$APPLICATION_DIR/Dockerfile.$FLAVOR" ]; then
      APPLICATION_DOCKERFILE="$APPLICATION_DIR/Dockerfile.$FLAVOR"
    fi
  fi

  prompt "Are you sure you want to build $APPLICATION_IMAGE from $APPLICATION_DOCKERFILE?"

  log "Building $APPLICATION_IMAGE image from $APPLICATION_DOCKERFILE"
  build_image \
    --build-arg="BASE_IMAGE=$BASE_IMAGE" \
    -t $APPLICATION_IMAGE \
    -f $APPLICATION_DOCKERFILE \
    $APPLICATION_DIR

  tag_image $APPLICATION_IMAGE
  push_image $APPLICATION_IMAGE

  exit 0
fi

# --------------------------------------------------------------------

prompt "Are you sure you want to build $BASE_IMAGE?"

log "Building base image: $BASE_IMAGE"

log "[STAGE 1]: Building neko-xorg-deps image"
build_image -t neko-xorg-deps runtime/xorg-deps/

log "[STAGE 2]: Building neko-server image"
build_image -t neko-server server/

log "[STAGE 3]: Building neko-client image"
build_image -t neko-client client/

if [ -z "$FLAVOR" ]; then
  RUNTIME_IMAGE="neko-runtime"
  log "[STAGE 4]: Building $RUNTIME_IMAGE image"
  build_image -t $RUNTIME_IMAGE runtime/
else
  RUNTIME_IMAGE="neko-$FLAVOR-runtime"
  log "[STAGE 4]: Building $RUNTIME_IMAGE image"
  build_image -t $RUNTIME_IMAGE -f runtime/Dockerfile.$FLAVOR runtime/
fi

log "[STAGE 5]: Building $BASE_IMAGE image"
build_image -t $BASE_IMAGE -f - . <<EOF
  FROM neko-xorg-deps AS xorg-deps
  FROM neko-server AS server
  FROM neko-client AS client
  FROM $RUNTIME_IMAGE AS runtime

  COPY --from=xorg-deps /usr/local/lib/xorg/modules/drivers/dummy_drv.so /usr/lib/xorg/modules/drivers/dummy_drv.so
  COPY --from=xorg-deps /usr/local/lib/xorg/modules/input/neko_drv.so /usr/lib/xorg/modules/input/neko_drv.so
  COPY --from=server /src/bin/plugins/ /etc/neko/plugins/
  COPY --from=server /src/bin/neko /usr/bin/neko
  COPY --from=client /src/dist/ /var/www
  
  COPY config.yml /etc/neko/neko.yaml
EOF

tag_image $BASE_IMAGE
push_image $BASE_IMAGE