#!/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
  python              Run Python extension tests
  typecheck           Run Python stub/typecheck/wheel verification
  lint-base           Run cargo clippy --no-deps -- -D warnings
  lint-serde          Run cargo clippy -F serde --no-deps -- -D warnings
  lint-strum          Run cargo clippy -F strum --no-deps -- -D warnings
  lint-pyo3           Run cargo clippy -F pyo3 --no-deps -- -D warnings
  test-base           Run cargo test --workspace --verbose
  test-serde          Run cargo test -F serde --workspace
  test-strum          Run cargo test -F strum --workspace
  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"
    cargo clippy --no-deps -- -D warnings
}

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

run_lint_strum() {
    ensure_clippy
    log "Running clippy with strum"
    cargo clippy -F strum --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 --workspace --verbose"
    rustc --version
    cargo --version
    cargo test --workspace --verbose
}

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

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

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_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
            ;;
        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-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 "$@"