#!/usr/bin/env bash

set -euo pipefail

ACTION="${1:-}"
if [[ ! "$ACTION" =~ ^(build|test|validate|run|dry-run|finalize|status|all)$ ]]; then
  echo "Usage: bash $0 {build|test|validate|run|dry-run|finalize|status|all}" >&2
  exit 2
fi

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="${PROJECT_DIR:-$(cd -- "$SCRIPT_DIR/.." && pwd)}"
PYTHON="${PYTHON:-python3}"
BACKEND="${BACKEND:-MKL}"
BACKEND="${BACKEND^^}"
BUILD_DIR="${BUILD_DIR:-$PROJECT_DIR/build-${BACKEND,,}}"
ENGINE="${ENGINE:-$BUILD_DIR/spectra_reml}"
MANIFEST="${MANIFEST:-}"
OUTPUT="${OUTPUT:-}"
THREADS="${THREADS:-1}"
BLAS_THREADS="${BLAS_THREADS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-256}"
RESUME="${RESUME:-1}"
FORCE="${FORCE:-0}"
CLI="$PROJECT_DIR/python/spectra_reml.py"

die() {
  echo "ERROR: $*" >&2
  exit 1
}

[[ "$BACKEND" == "MKL" || "$BACKEND" == "OPENBLAS" ]] || die "BACKEND must be MKL or OPENBLAS"
[[ "$THREADS" =~ ^[1-9][0-9]*$ ]] || die "THREADS must be positive"
[[ "$BLAS_THREADS" =~ ^[1-9][0-9]*$ ]] || die "BLAS_THREADS must be positive"
[[ "$BLOCK_SIZE" =~ ^[1-9][0-9]*$ ]] || die "BLOCK_SIZE must be positive"
[[ "$RESUME" == "0" || "$RESUME" == "1" ]] || die "RESUME must be 0 or 1"
[[ "$FORCE" == "0" || "$FORCE" == "1" ]] || die "FORCE must be 0 or 1"
[[ "$RESUME" != "1" || "$FORCE" != "1" ]] || die "RESUME and FORCE are mutually exclusive"

build_project() {
  local args=(
    -S "$PROJECT_DIR"
    -B "$BUILD_DIR"
    -DCMAKE_BUILD_TYPE=Release
    -DREML_BLAS="$BACKEND"
    -DREML_ENABLE_OPENMP=ON
  )
  if [[ "$BACKEND" == "MKL" ]]; then
    args+=(
      -DMKL_INTERFACE=lp64
      -DMKL_LINK=dynamic
      -DMKL_THREADING=sequential
    )
  fi
  cmake "${args[@]}"
  cmake --build "$BUILD_DIR" --parallel
}

if [[ "$ACTION" == "build" ]]; then
  build_project
  exit 0
fi

command -v "$PYTHON" >/dev/null 2>&1 || die "Python was not found: $PYTHON"
[[ -s "$CLI" ]] || die "Generic CLI is missing: $CLI"

if [[ "$ACTION" == "test" ]]; then
  [[ -d "$BUILD_DIR" ]] || die "Build first: bash $0 build"
  ctest --test-dir "$BUILD_DIR" --output-on-failure
  "$PYTHON" "$PROJECT_DIR/tests/test_math_reference.py"
  "$PYTHON" "$PROJECT_DIR/tests/python/test_spectra_reml_cli.py"
  exit 0
fi

[[ -n "$MANIFEST" ]] || die "Set MANIFEST=/absolute/path/to/manifest.json"

case "$ACTION" in
  validate|status)
    "$PYTHON" "$CLI" "$ACTION" --manifest "$MANIFEST"
    ;;
  finalize)
    [[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
    "$PYTHON" "$CLI" finalize --manifest "$MANIFEST" --output "$OUTPUT"
    ;;
  run|dry-run|all)
    if [[ "$ACTION" != "dry-run" ]]; then
      [[ -x "$ENGINE" ]] || die "Engine is missing/not executable: $ENGINE"
    fi
    mode=()
    if [[ "$FORCE" == "1" ]]; then
      mode+=(--force)
    elif [[ "$RESUME" == "1" ]]; then
      mode+=(--resume)
    fi
    common=(
      --manifest "$MANIFEST"
      --engine "$ENGINE"
      --threads "$THREADS"
      --blas-threads "$BLAS_THREADS"
      --block-size "$BLOCK_SIZE"
      "${mode[@]}"
    )
    if [[ "$ACTION" == "dry-run" ]]; then
      "$PYTHON" "$CLI" run "${common[@]}" --dry-run
    elif [[ "$ACTION" == "all" ]]; then
      [[ -n "$OUTPUT" ]] || die "Set OUTPUT=/absolute/path/to/results.tsv.gz"
      "$PYTHON" "$CLI" all "${common[@]}" --output "$OUTPUT"
    else
      "$PYTHON" "$CLI" run "${common[@]}"
    fi
    ;;
esac