#!/usr/bin/env bash

set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="${VENV_DIR:-$ROOT_DIR/.venv}"

usage() {
    cat <<'EOF'
Usage: ./test.sh [suite...]

Suites:
  all                 Run every suite from .gitlab-ci.yml (default)
  lint                Run all clippy jobs
  rust                Run all Rust test jobs
  plain               Run all plain storage test jobs
  python              Run Python extension tests
  typecheck           Run Python stub/typecheck/wheel verification
  lint-base           Run cargo clippy -F ndarray --no-deps -- -D warnings
  lint-serde          Run cargo clippy -F "serde ndarray" --no-deps -- -D warnings
  lint-strum          Run cargo clippy -F "strum ndarray" --no-deps -- -D warnings
  lint-pyo3           Run cargo clippy -F pyo3 --no-deps -- -D warnings
  test-base           Run cargo test -F ndarray --workspace --verbose
  test-serde          Run cargo test -F "serde ndarray" --workspace
  test-strum          Run cargo test -F "strum ndarray" --workspace
  test-no-ndarray     Run cargo test --no-default-features --workspace
  test-serde-no-ndarray
                      Run cargo test --no-default-features -F serde --workspace
  test-strum-no-ndarray
                      Run cargo test --no-default-features -F strum --workspace
  test-plain-f64      Run cargo test --no-default-features -F storage-f64 --workspace
  test-plain-f64-serde
                      Run cargo test --no-default-features -F "storage-f64 serde" --workspace
  test-plain-f32      Run cargo test --no-default-features -F storage-f32 --workspace
  test-plain-f32-serde
                      Run cargo test --no-default-features -F "storage-f32 serde" --workspace
  check-nostd-f64     Run cargo check --no-default-features -F "no_std storage-f64"
  check-nostd-f32     Run cargo check --no-default-features -F "no_std storage-f32"
  test-pyo3           Run pytest against a local maturin develop install
  typecheck-pyo3      Run mypy and wheel typing artifact verification

Environment:
  VENV_DIR            Python virtualenv directory (default: ./.venv)
  SKIP_PIP_INSTALL=1  Reuse an existing virtualenv without installing packages
EOF
}

log() {
    printf '\n[%s] %s\n' "$(date '+%H:%M:%S')" "$*"
}

require_cmd() {
    if ! command -v "$1" >/dev/null 2>&1; then
        echo "Missing required command: $1" >&2
        exit 1
    fi
}

ensure_clippy() {
    if ! rustup component list --installed | grep -q '^clippy'; then
        log "Installing clippy"
        rustup component add clippy
    fi
}

ensure_python_venv() {
    require_cmd python3
    if [[ ! -d "$VENV_DIR" ]]; then
        log "Creating virtualenv at $VENV_DIR"
        python3 -m venv "$VENV_DIR"
    fi

    # shellcheck disable=SC1090
    source "$VENV_DIR/bin/activate"

    if [[ "${SKIP_PIP_INSTALL:-0}" != "1" ]]; then
        log "Installing Python test dependencies"
        python -m pip install --upgrade pip setuptools wheel maturin pytest numpy "mypy==1.8.0"
    fi
}

run_lint_base() {
    ensure_clippy
    log "Running clippy with ndarray"
    cargo clippy -F ndarray --no-deps -- -D warnings
}

run_lint_serde() {
    ensure_clippy
    log "Running clippy with serde and ndarray"
    cargo clippy -F "serde ndarray" --no-deps -- -D warnings
}

run_lint_strum() {
    ensure_clippy
    log "Running clippy with strum and ndarray"
    cargo clippy -F "strum ndarray" --no-deps -- -D warnings
}

run_lint_pyo3() {
    ensure_clippy
    log "Running clippy with pyo3"
    cargo clippy -F pyo3 --no-deps -- -D warnings
}

run_test_base() {
    log "Running cargo test -F ndarray --workspace --verbose"
    rustc --version
    cargo --version
    cargo test -F ndarray --workspace --verbose
}

run_test_serde() {
    log "Running cargo test -F \"serde ndarray\" --workspace"
    rustc --version
    cargo --version
    cargo test -F "serde ndarray" --workspace
}

run_test_strum() {
    log "Running cargo test -F \"strum ndarray\" --workspace"
    cargo test -F "strum ndarray" --workspace
}

run_test_no_ndarray() {
    log "Running cargo test --no-default-features --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features --workspace
}

run_test_serde_no_ndarray() {
    log "Running cargo test --no-default-features -F serde --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features -F serde --workspace
}

run_test_strum_no_ndarray() {
    log "Running cargo test --no-default-features -F strum --workspace"
    cargo test --no-default-features -F strum --workspace
}

run_test_plain_f64() {
    log "Running cargo test --no-default-features -F storage-f64 --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features -F storage-f64 --workspace
}

run_test_plain_f64_serde() {
    log "Running cargo test --no-default-features -F \"storage-f64 serde\" --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features -F "storage-f64 serde" --workspace
}

run_test_plain_f32() {
    log "Running cargo test --no-default-features -F storage-f32 --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features -F storage-f32 --workspace
}

run_test_plain_f32_serde() {
    log "Running cargo test --no-default-features -F \"storage-f32 serde\" --workspace"
    rustc --version
    cargo --version
    cargo test --no-default-features -F "storage-f32 serde" --workspace
}

run_test_plain() {
    run_test_plain_f64
    run_test_plain_f64_serde
    run_test_plain_f32
    run_test_plain_f32_serde
}

run_check_nostd_f64() {
    log "Running cargo check --no-default-features -F \"no_std storage-f64\""
    rustc --version
    cargo --version
    cargo check --no-default-features -F "no_std storage-f64"
}

run_check_nostd_f32() {
    log "Running cargo check --no-default-features -F \"no_std storage-f32\""
    rustc --version
    cargo --version
    cargo check --no-default-features -F "no_std storage-f32"
}

run_test_pyo3() {
    ensure_python_venv
    log "Generating Python stubs"
    python scripts/generate_python_stubs.py
    log "Installing extension into virtualenv with maturin"
    maturin develop --features pyo3
    log "Running pytest"
    pytest tests
    deactivate
}

run_typecheck_pyo3() {
    ensure_python_venv
    log "Generating Python stubs"
    python scripts/generate_python_stubs.py
    log "Checking generated Python stubs are up to date"
    python scripts/generate_python_stubs.py --check
    log "Running mypy against source package"
    MYPYPATH=python mypy --python-version 3.7 tests/typecheck/mypy_smoke.py
    log "Building wheel"
    rm -rf dist
    maturin build --features pyo3 --out dist
    log "Verifying wheel typing artifacts"
    python - <<'PY'
import glob
import zipfile

wheels = sorted(glob.glob("dist/*.whl"))
if not wheels:
    raise SystemExit("No wheels built in dist/")

required = {
    "unitforge/__init__.py",
    "unitforge/__init__.pyi",
    "unitforge/py.typed",
}
with zipfile.ZipFile(wheels[-1]) as wheel:
    names = set(wheel.namelist())
missing = sorted(required - names)
if missing:
    raise SystemExit(f"Wheel is missing required typing artifacts: {missing}")
print("Wheel typing artifacts verified.")
PY
    log "Installing built wheel"
    python -m pip install --force-reinstall --no-deps dist/*.whl
    local tmp_dir
    tmp_dir="$(mktemp -d)"
    cp tests/typecheck/mypy_smoke.py "$tmp_dir"/
    log "Running mypy against installed wheel from isolated directory"
    (
        cd "$tmp_dir"
        mypy --python-version 3.7 mypy_smoke.py
    )
    rm -rf "$tmp_dir"
    deactivate
}

run_suite() {
    case "$1" in
        all)
            run_lint_base
            run_lint_serde
            run_lint_strum
            run_lint_pyo3
            run_test_base
            run_test_serde
            run_test_strum
            run_test_no_ndarray
            run_test_serde_no_ndarray
            run_test_strum_no_ndarray
            run_test_plain
            run_check_nostd_f64
            run_check_nostd_f32
            run_test_pyo3
            run_typecheck_pyo3
            ;;
        lint)
            run_lint_base
            run_lint_serde
            run_lint_strum
            run_lint_pyo3
            ;;
        rust)
            run_test_base
            run_test_serde
            run_test_strum
            run_test_no_ndarray
            run_test_serde_no_ndarray
            run_test_strum_no_ndarray
            run_test_plain
            run_check_nostd_f64
            run_check_nostd_f32
            ;;
        plain)
            run_test_plain
            ;;
        python)
            run_test_pyo3
            ;;
        typecheck)
            run_typecheck_pyo3
            ;;
        lint-base) run_lint_base ;;
        lint-serde) run_lint_serde ;;
        lint-strum) run_lint_strum ;;
        lint-pyo3) run_lint_pyo3 ;;
        test-base) run_test_base ;;
        test-serde) run_test_serde ;;
        test-strum) run_test_strum ;;
        test-no-ndarray) run_test_no_ndarray ;;
        test-serde-no-ndarray) run_test_serde_no_ndarray ;;
        test-strum-no-ndarray) run_test_strum_no_ndarray ;;
        test-plain-f64) run_test_plain_f64 ;;
        test-plain-f64-serde) run_test_plain_f64_serde ;;
        test-plain-f32) run_test_plain_f32 ;;
        test-plain-f32-serde) run_test_plain_f32_serde ;;
        check-nostd-f64) run_check_nostd_f64 ;;
        check-nostd-f32) run_check_nostd_f32 ;;
        test-pyo3) run_test_pyo3 ;;
        typecheck-pyo3) run_typecheck_pyo3 ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown suite: $1" >&2
            usage >&2
            exit 1
            ;;
    esac
}

main() {
    require_cmd cargo
    require_cmd rustc

    cd "$ROOT_DIR"

    if [[ $# -eq 0 ]]; then
        run_suite all
        return
    fi

    for suite in "$@"; do
        run_suite "$suite"
    done
}

main "$@"