#!/usr/bin/env bash

set -euo pipefail

toplevel="$(git rev-parse --show-toplevel)"
DEFAULT_COMPLEMENT_SRC="$toplevel/target/complement"
DEFAULT_COMPLEMENT_REVISION="6d2fdc286c2b44faaddd1037205869b2242a4005"

MSC="${COMPLEMENT_MSC:-}"
IMAGE_ARCHIVE=""
while [ "$#" -gt 0 ]; do
    case "$1" in
        --msc)
            if [ "$#" -lt 2 ]; then
                echo "--msc requires an MSC number"
                exit 2
            fi
            MSC="$2"
            shift 2
            ;;
        --image)
            if [ "$#" -lt 2 ]; then
                echo "--image requires a path to a Docker image archive"
                exit 2
            fi
            IMAGE_ARCHIVE="$2"
            shift 2
            ;;
        *)
            break
            ;;
    esac
done

if [ -n "$MSC" ]; then
    MSC="${MSC#MSC}"
    MSC="${MSC#msc}"
    if ! [[ "$MSC" =~ ^[0-9]+$ ]]; then
        echo "--msc must be an MSC number, for example: --msc 3391"
        exit 2
    fi
fi

# The root path where Complement is available.
if [ -n "${COMPLEMENT_SRC:-}" ]; then
    LOG_FILE="${1:-tests/test_results/complement/test_logs.jsonl}"
    RESULTS_FILE="${2:-tests/test_results/complement/test_results.jsonl}"
else
    COMPLEMENT_SRC="${1:-$DEFAULT_COMPLEMENT_SRC}"
    LOG_FILE="${2:-tests/test_results/complement/test_logs.jsonl}"
    RESULTS_FILE="${3:-tests/test_results/complement/test_results.jsonl}"
fi

# The base docker image to use for complement tests
# You can build the default with `docker build -t continuwuity:complement -f ./docker/complement.Dockerfile .`
# after running `cargo build`. Only the debug binary is used.
COMPLEMENT_BASE_IMAGE="${COMPLEMENT_BASE_IMAGE:-continuwuity:complement}"

# Complement tests that are skipped due to flakiness/reliability issues or we don't implement such features and won't for a long time
SKIPPED_COMPLEMENT_TESTS='TestPartialStateJoin.*|TestRoomDeleteAlias/Parallel/Regular_users_can_add_and_delete_aliases_when_m.*|TestRoomDeleteAlias/Parallel/Can_delete_canonical_alias|TestUnbanViaInvite.*|TestRoomState/Parallel/GET_/publicRooms_lists.*"|TestRoomDeleteAlias/Parallel/Users_with_sufficient_power-level_can_delete_other.*'

# Download the pinned upstream Complement checkout on first use. A local
# checkout can be supplied as the first argument or with COMPLEMENT_SRC.
if [ ! -e "$COMPLEMENT_SRC" ]; then
    mkdir -p "$(dirname "$COMPLEMENT_SRC")"
    git init --quiet "$COMPLEMENT_SRC"
    git -C "$COMPLEMENT_SRC" remote add origin https://github.com/matrix-org/complement.git
    git -C "$COMPLEMENT_SRC" fetch --depth 1 origin "$DEFAULT_COMPLEMENT_REVISION"
    git -C "$COMPLEMENT_SRC" checkout --quiet --detach FETCH_HEAD
fi

if [ ! -d "$COMPLEMENT_SRC" ] || [ ! -f "$COMPLEMENT_SRC/go.mod" ]; then
    echo "\$COMPLEMENT_SRC must be a directory containing Complement source code"
    exit 1
fi

pushd "$toplevel" > /dev/null

for command in docker go jq; do
    if ! command -v "$command" > /dev/null; then
        echo "missing required command: $command"
        exit 1
    fi
done

if [ -n "$IMAGE_ARCHIVE" ]; then
    if [ ! -f "$IMAGE_ARCHIVE" ]; then
        echo "Docker image archive not found: $IMAGE_ARCHIVE"
        exit 1
    fi
    load_output="$(docker load --input "$IMAGE_ARCHIVE" 2>&1)"
    printf '%s\n' "$load_output"
    COMPLEMENT_BASE_IMAGE="$(printf '%s\n' "$load_output" | sed -n 's/^Loaded image: //p; s/^Loaded image ID: //p' | tail -n1)"
    if [ -z "$COMPLEMENT_BASE_IMAGE" ]; then
        echo "could not determine the image loaded from: $IMAGE_ARCHIVE"
        exit 1
    fi
fi

if ! docker image inspect "$COMPLEMENT_BASE_IMAGE" > /dev/null 2>&1; then
    echo "Complement base image not found: $COMPLEMENT_BASE_IMAGE"
    echo "Build it with: cargo build -p conduwuit && docker build -t $COMPLEMENT_BASE_IMAGE -f docker/complement.Dockerfile ."
    exit 1
fi

mkdir -p "$(dirname "$LOG_FILE")" "$(dirname "$RESULTS_FILE")"
LOG_TMP="$(mktemp "${LOG_FILE}.XXXXXX")"
RESULTS_TMP="$(mktemp "${RESULTS_FILE}.XXXXXX")"
trap 'rm -f "$LOG_TMP" "$RESULTS_TMP"' EXIT

echo ""
echo "running go test with:"
echo "\$COMPLEMENT_SRC: $COMPLEMENT_SRC"
COMPLEMENT_REVISION="$(git -C "$COMPLEMENT_SRC" rev-parse HEAD 2>/dev/null || echo unknown)"
echo "\$COMPLEMENT_REVISION: $COMPLEMENT_REVISION"
echo "\$COMPLEMENT_MSC: ${MSC:-all}"
echo "\$COMPLEMENT_BASE_IMAGE: $COMPLEMENT_BASE_IMAGE"
echo "\$RESULTS_FILE: $RESULTS_FILE"
echo "\$LOG_FILE: $LOG_FILE"
echo ""

TEST_PACKAGES=(./tests/...)
if [ -n "$MSC" ]; then
    test_files="$({
        grep -rilE --include='*.go' "msc$MSC([^0-9]|$)" "$COMPLEMENT_SRC/tests" || true
        find "$COMPLEMENT_SRC/tests" -type f -name '*.go' -print | while IFS= read -r file; do
            if [[ "$(basename "$file")" =~ [Mm][Ss][Cc]$MSC([^0-9]|$) ]]; then
                printf '%s\n' "$file"
            fi
        done
    } | sort -u)"
    if [ -z "$test_files" ]; then
        echo "no Complement tests mention MSC$MSC"
        exit 1
    fi
    TEST_PACKAGES=()
    while IFS= read -r test_file; do
        test_dir="$(dirname "$test_file")"
        TEST_PACKAGES+=(".${test_dir#"$COMPLEMENT_SRC"}")
    done <<< "$test_files"
fi

# `COMPLEMENT_ENABLE_DIRTY_RUNS=1`  reuses the same complement container for faster complement, at the possible expense of test environment pollution
set +e
(cd "$COMPLEMENT_SRC" && \
    COMPLEMENT_BASE_IMAGE="$COMPLEMENT_BASE_IMAGE" \
    go test -tags="conduwuit_blacklist" -skip="$SKIPPED_COMPLEMENT_TESTS" -v -timeout 1h -json "${TEST_PACKAGES[@]}"
) | tee "$LOG_TMP"
test_status=${PIPESTATUS[0]}
set -e

# Post-process the results into an easy-to-compare format, sorted by Test name for reproducible results
jq -s -c 'sort_by(.Test)[]' < "$LOG_TMP" | jq -c '
    select(
        (.Action == "pass" or .Action == "fail" or .Action == "skip")
        and .Test != null
    ) | {Action: .Action, Test: .Test}
    ' > "$RESULTS_TMP"

mv "$LOG_TMP" "$LOG_FILE"
mv "$RESULTS_TMP" "$RESULTS_FILE"
trap - EXIT

echo ""
echo ""
echo "complement logs saved at $LOG_FILE"
echo "complement results saved at $RESULTS_FILE"
echo ""
echo ""

exit "$test_status"