#!/usr/bin/env bash

set -e

if [ "${SW_TASK_DISABLE_DEBUG}" != "1" ]; then
    set -x
fi

ulimit -n 65535 || true

VERBOSE="-vvv"
HANDLER=${SW_RUN_HANDLER:-""}
STEP=${SW_TASK_STEP:-""}
TASK_INDEX=${SW_TASK_INDEX:-0}
TASK_NUM=${SW_TASK_NUM:-0}
RUNTIME_RESTORED=${SW_USER_RUNTIME_RESTORED:-0}
RUNTIME_WORKDIR=${SW_USER_RUNTIME_WORKDIR:-/opt/starwhale.user/runtime}


welcome() {
    echo "===================================="
    echo "StarWhale Docker Entrypoint"
    echo "Date: `date -u +%Y-%m-%dT%H:%M:%SZ`"
    echo "Version: `swcli --version`"
    echo "Run: $1 "
    echo "Instance: ${SW_INSTANCE_URI}"
    echo "Project: ${SW_PROJECT_URI}, ${SW_PROJECT}"
    echo "Model Version: ${SW_MODEL_URI}, ${SW_MODEL_VERSION}"
    echo "Runtime Version: ${SW_RUNTIME_URI}, ${SW_RUNTIME_VERSION}"
    echo "Runtime Restored: ${SW_USER_RUNTIME_RESTORED}"
    echo "Local User: ${SW_USER:-root}"
    echo "===================================="
    if [ ! -z "${SW_USER}" ];
    then
      useradd -g ${SW_USER_GROUP_ID} -u ${SW_USER_ID} ${SW_USER}
      su ${SW_USER}
    fi
}

prepare(){
    if [ "${RUNTIME_RESTORED}" == "1" ]; then
        echo '-->[Preparing] Runtime has been restored, activate it...'
        source "${RUNTIME_WORKDIR}"/env.sw
        $(bash "${RUNTIME_WORKDIR}"/activate.sw)
    fi

    if [ "${SW_INSTANCE_URI}" != "local" ]; then
        echo '-->[Preparing] login to server...'
        swcli instance login --token "${SW_TOKEN}" --alias server "${SW_INSTANCE_URI}"
    fi
}

# since 0.6.0
compatibility_check() {
    if [ -z "${SW_PROJECT_URI}" ] || [ -z "${SW_MODEL_URI}" ] || { [ -z "${SW_RUNTIME_URI}" ] && [ "${RUNTIME_RESTORED}" != "1" ];}; then
        echo "-->[Error] The following variables are missing: SW_PROJECT_URI or SW_MODEL_URI or SW_RUNTIME_URI, possible reason: your server version is lower, please upgrade the server to at least 0.6.0, now will exit."
        exit 1
    fi;
}

run() {
    if [ "${SW_TASK_DISABLE_DEBUG}" = "1" ]; then
        echo "-->[Preparing] debug config ..."
        VERBOSE="-v"
    fi
    echo "-->[Running] start to run model: ${STEP}, use $(which swcli) cli ..."
    if [ "${RUNTIME_RESTORED}" != "1" ]; then
        runtime_args="--runtime ${SW_RUNTIME_URI}"
    else
        runtime_args="--forbid-packaged-runtime"
    fi

    swcli ${VERBOSE} model run --handler="${HANDLER}" --step="${STEP}" \
        --task-index="${TASK_INDEX}" --override-task-num="${TASK_NUM}" \
        --uri="${SW_MODEL_URI}" --version="${SW_JOB_VERSION}" \
        --log-project="${SW_PROJECT_URI}" \
        --forbid-snapshot \
        ${runtime_args} -- ${SW_TASK_EXTRA_CMD_ARGS} || exit 1
}

serve() {
    echo "-->[Serving] start to serve, use $(which swcli) cli ..."
    export
    if [ "${RUNTIME_RESTORED}" != "1" ]; then
        swcli ${VERBOSE} model serve --uri "${SW_MODEL_URI}" --runtime "${SW_RUNTIME_URI}" --host 0.0.0.0 || exit 1
    else
        swcli ${VERBOSE} model serve --uri "${SW_MODEL_URI}" --host 0.0.0.0 --forbid-packaged-runtime || exit 1
    fi
}

run_code_server () {
    echo "-->[Preparing] run code-server ..."
    if command -v code-server &> /dev/null
    then
        PASSWORD=${SW_DEV_TOKEN:=$SW_JOB_VERSION} nohup code-server --bind-addr 0.0.0.0:${SW_DEV_PORT:-8000} --disable-telemetry --disable-update-check --user-data-dir /tmp/code-server-data > /var/log/dev.log 2>&1 &
        echo "-->[Prepared] run code-server done."
    else
        echo "code-server is not installed, now will exit." && exit 1
    fi
}

ds_build_and_upload () {
    if [ -z "${SW_PROJECT_URI}" ]; then
        echo "-->[Error] There needs an environment variable: SW_PROJECT_URI, now will exit."
        exit 1
    fi;
    if [ -n "$DATASET_BUILD_DIR_PREFIX" ]; then
      echo "-->[Preparing] Downloading files..."
      BUILD_DIR="$DATASET_BUILD_NAME"
      mkdir -p "$BUILD_DIR"
      cd "$BUILD_DIR"

      SIGNED_URLS=$(curl -X 'GET' "$SW_INSTANCE_URI/api/v1/filestorage/signedurl?pathPrefix=$DATASET_BUILD_DIR_PREFIX"  -H 'accept: application/json'  -H "Authorization: $SW_TOKEN" | jq  ".data.signedUrls")

      for entry in $(echo "$SIGNED_URLS" | jq -r 'to_entries|map("\(.key)@\(.value)")|.[]'); do
        IFS='@' read -r file signedurl <<< "$entry"

        filedir=$(dirname "$file")
        if [ ! -d "$filedir" ]; then
          mkdir -p "$filedir"
        fi
        echo "$file $signedurl"
      done | xargs -I {} -n 2 -P 10 sh -c 'curl -o "$1" "$2"' sh
      cd -
    fi

    cmd="swcli dataset build -n '$DATASET_BUILD_NAME'"
    if [ "$DATASET_BUILD_TYPE" = "IMAGE" ]; then
      cmd="$cmd -if $BUILD_DIR"
    elif [ "$DATASET_BUILD_TYPE" = "VIDEO" ]; then
      cmd="$cmd -vf $BUILD_DIR"
    elif [ "$DATASET_BUILD_TYPE" = "AUDIO" ]; then
      cmd="$cmd -af $BUILD_DIR"
    elif [ "$DATASET_BUILD_TYPE" = "CSV" ]; then
      cmd="$cmd -c $BUILD_DIR"
      if [ -n "$DATASET_BUILD_CSV_DIALECT" ]; then
        cmd="$cmd --dialect '$DATASET_BUILD_CSV_DIALECT'"
      fi
      if [ -n "$DATASET_BUILD_CSV_DELIMITER" ]; then
        cmd="$cmd --delimiter '$DATASET_BUILD_CSV_DELIMITER'"
      fi
      if [ -n "$DATASET_BUILD_CSV_QUOTECHAR" ]; then
        cmd="$cmd --quotechar '$DATASET_BUILD_CSV_QUOTECHAR'"
      fi
    elif [ "$DATASET_BUILD_TYPE" = "JSON" ]; then
      cmd="$cmd -jf $BUILD_DIR"
      if [ -n "$DATASET_BUILD_JSON_FIELD_SELECTOR" ]; then
        cmd="$cmd --field-selector '$DATASET_BUILD_JSON_FIELD_SELECTOR'"
      fi
    elif [ "$DATASET_BUILD_TYPE" = "HUGGINGFACE" ]; then
      if [ -z "$DATASET_BUILD_HF_REPO" ]; then
        echo "-->[Error] Build dataset by HuggingFace required env: DATASET_BUILD_HF_REPO."
        exit 1
      fi
      cmd="$cmd -hf '$DATASET_BUILD_HF_REPO'"
      if [ -n "$DATASET_BUILD_HF_SUBSET" ]; then
        cmd="$cmd --subset '$DATASET_BUILD_HF_SUBSET'"
      fi
      if [ -n "$DATASET_BUILD_HF_SPLIT" ]; then
        cmd="$cmd --split '$DATASET_BUILD_HF_SPLIT'"
      fi
      if [ -n "$DATASET_BUILD_HF_REVISION" ]; then
        cmd="$cmd --revision '$DATASET_BUILD_HF_REVISION'"
      fi
    else
      echo "Unknown type: $DATASET_BUILD_TYPE" && exit 1
    fi

    echo "-->[Building] Start to build dataset: $DATASET_BUILD_NAME..."
    eval "$cmd" || exit 1

    echo "-->[Uploading] Start to upload dataset: $DATASET_BUILD_NAME..."
    swcli instance login --token "$SW_TOKEN" --alias server "$SW_INSTANCE_URI"
    swcli dataset copy --patch "$DATASET_BUILD_NAME"/version/latest "${SW_PROJECT_URI}" || exit 1
}

welcome "$1"
case "$1" in
    run|evaluation|fine_tune)
        compatibility_check && prepare && run
        ;;
    serve|serving)
        compatibility_check && prepare && serve
        ;;
    dev)
        prepare && run_code_server
        tail -f /var/log/dev.log
        ;;
    dataset_build)
        ds_build_and_upload
        ;;
    *)
        prepare "starwhale" && exec "$@"
        ;;
esac
