#!/usr/bin/env bash
set -euo pipefail

CONFIG_PATH="config.yaml"
CALIBRATION_PATH="examples/training/data/aihub_calibration"
INPUT_FILE="examples/training/data/inputs.npz"
FROM_STEP="quantize"
FROM_JOB=""
MODEL_S3_URI=""
ONNX_PATH=""
INPUT_NAME=""
DOWNLOAD=true
OUTPUT_PATH=""

usage() {
  cat <<EOF
Usage: $0 [options]

Options:
  --config PATH              Path to qc-cli config file. Default: config.yaml
  --calibration PATH         Calibration .npz file or directory of .npy samples.
                              Default: ${CALIBRATION_PATH}
  --input-file PATH          Validation .npz or .npy inputs. Default: ${INPUT_FILE}
  --from-step STEP           Resume upload from: quantize, compile, validate, profile.
                              Default: ${FROM_STEP}
  --from-job NAME            SageMaker training job whose model artifact should upload.
                              Defaults to the last training job in local qc-cli state.
  --model-s3-uri URI         S3 URI of model.tar.gz to upload.
  --onnx-path PATH           Local ONNX path or ONNX path inside extracted artifact.
  --input-name NAME          Input name for .npy validation files.
  --skip-download            Do not download the compiled AI Hub artifact after upload.
  --output PATH              Destination file for ai-hub download.
  -h, --help                 Show this help.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --config)
      CONFIG_PATH="$2"
      shift 2
      ;;
    --calibration)
      CALIBRATION_PATH="$2"
      shift 2
      ;;
    --input-file)
      INPUT_FILE="$2"
      shift 2
      ;;
    --from-step)
      FROM_STEP="$2"
      shift 2
      ;;
    --from-job)
      FROM_JOB="$2"
      shift 2
      ;;
    --model-s3-uri)
      MODEL_S3_URI="$2"
      shift 2
      ;;
    --onnx-path)
      ONNX_PATH="$2"
      shift 2
      ;;
    --input-name)
      INPUT_NAME="$2"
      shift 2
      ;;
    --skip-download)
      DOWNLOAD=false
      shift
      ;;
    --output)
      OUTPUT_PATH="$2"
      shift 2
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
done

if [[ ! -f "${CONFIG_PATH}" ]]; then
  echo "Config not found: ${CONFIG_PATH}" >&2
  exit 1
fi

case "${FROM_STEP}" in
  quantize|compile|validate|profile)
    ;;
  *)
    echo "--from-step must be one of: quantize, compile, validate, profile" >&2
    exit 1
    ;;
esac

if [[ ! -e "${CALIBRATION_PATH}" ]]; then
  echo "Calibration path not found: ${CALIBRATION_PATH}" >&2
  echo "Pass --calibration with a .npz file or directory of .npy samples." >&2
  exit 1
fi

if [[ ! -f "${INPUT_FILE}" ]]; then
  echo "Input file not found: ${INPUT_FILE}" >&2
  echo "Pass --input-file with a validation .npz or .npy file." >&2
  exit 1
fi

run() {
  echo "+ $*"
  "$@"
}

UPLOAD_ARGS=(
  "${CALIBRATION_PATH}"
  "${INPUT_FILE}"
  --from-step "${FROM_STEP}"
  --config "${CONFIG_PATH}"
)

if [[ -n "${FROM_JOB}" ]]; then
  UPLOAD_ARGS+=(--from-job "${FROM_JOB}")
fi

if [[ -n "${MODEL_S3_URI}" ]]; then
  UPLOAD_ARGS+=(--model-s3-uri "${MODEL_S3_URI}")
fi

if [[ -n "${ONNX_PATH}" ]]; then
  UPLOAD_ARGS+=(--onnx-path "${ONNX_PATH}")
fi

if [[ -n "${INPUT_NAME}" ]]; then
  UPLOAD_ARGS+=(--input-name "${INPUT_NAME}")
fi

run uv run qc-cli ai-hub upload "${UPLOAD_ARGS[@]}"

if [[ "${DOWNLOAD}" == false ]]; then
  exit 0
fi

DOWNLOAD_ARGS=(--config "${CONFIG_PATH}")
if [[ -n "${OUTPUT_PATH}" ]]; then
  DOWNLOAD_ARGS+=(--output "${OUTPUT_PATH}")
fi

run uv run qc-cli ai-hub download "${DOWNLOAD_ARGS[@]}"