#!/usr/bin/env bash
# Run the test suite for explainer_env.
#
# Usage:
#   tests/run_tests.sh              # fast tests (models, task_bank, rewards, environment)
#   tests/run_tests.sh --all        # fast + client-server integration
#   tests/run_tests.sh --docker     # fast + docker build & test
#   tests/run_tests.sh --full       # everything

set -euo pipefail
cd "$(dirname "$0")/.."  # explainer_env/

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m'

PASSED=0
FAILED=0
SKIPPED=0

run() {
    local label="$1"; shift
    printf "%-40s" "$label"
    if output=$("$@" 2>&1); then
        echo -e "${GREEN}OK${NC}"
        PASSED=$((PASSED + 1))
    else
        echo -e "${RED}FAIL${NC}"
        echo "$output" | tail -5
        FAILED=$((FAILED + 1))
    fi
}

skip() {
    printf "%-40s" "$1"
    echo -e "${YELLOW}SKIP${NC}"
    SKIPPED=$((SKIPPED + 1))
}

echo "=== explainer_env test suite ==="
echo ""

# --- Fast tests (no server needed) ---
echo "--- Unit tests ---"
run "models"      uv run python tests/test_models.py
run "task_bank"   uv run python tests/test_task_bank.py
run "rewards"     uv run python tests/test_rewards.py
run "environment" uv run python tests/test_environment.py
run "ruff lint"   uvx ruff check .

# --- Integration tests (need server / docker) ---
MODE="${1:-}"

if [[ "$MODE" == "--all" || "$MODE" == "--full" ]]; then
    echo ""
    echo "--- Client-server integration ---"
    run "client_server" uv run python tests/test_client_server.py
else
    echo ""
    skip "client_server (use --all)"
fi

if [[ "$MODE" == "--docker" || "$MODE" == "--full" ]]; then
    echo ""
    echo "--- Docker integration ---"
    run "docker" uv run python tests/test_docker.py
else
    skip "docker (use --docker or --full)"
fi

# --- Summary ---
echo ""
TOTAL=$((PASSED + FAILED + SKIPPED))
echo "=== ${PASSED} passed, ${FAILED} failed, ${SKIPPED} skipped (${TOTAL} total) ==="
[[ $FAILED -eq 0 ]] || exit 1