#!/usr/bin/env bash
set -euo pipefail
# note: needs `sudo modprobe vhost_vsock`!

log() {
    printf "\n\033[1;36m>>> %s\033[0m\n" "$*"
}

strip_ansi() {
    local esc
    esc=$(printf '\033')
    sed -E "s/${esc}\[[0-9;]*[a-zA-Z]//g; s/${esc}\([a-zA-Z]//g" "$@"
}

declare -a TEST_NAMES=()
declare -a TEST_STATUSES=()
declare -a TEST_TIMES=()

get_time_ms() {
    local t="${EPOCHREALTIME:-}"
    if [[ "$t" == *.* ]]; then
        local secs="${t%.*}"
        local subs="${t#*.}"
        subs="${subs:0:3}"
        while [ "${#subs}" -lt 3 ]; do
            subs="${subs}0"
        done
        echo "${secs}${subs}"
    else
        echo "$(date +%s)000"
    fi
}

format_duration() {
    local ms=$1
    local secs=$((ms / 1000))
    local rem=$((ms % 1000))
    printf "%d.%03ds" "$secs" "$rem"
}

print_summary() {
    if [ "${#TEST_NAMES[@]}" -eq 0 ]; then
        return
    fi
    printf "\n"
    log "test summary"
    echo "========================================="
    local passed_count=0
    local failed_count=0
    local total_time=0
    for i in "${!TEST_NAMES[@]}"; do
        local name="${TEST_NAMES[$i]}"
        local status="${TEST_STATUSES[$i]}"
        local duration_ms="${TEST_TIMES[$i]}"
        local duration_str
        duration_str=$(format_duration "$duration_ms")

        local status_color="\033[0;32m"
        if [ "$status" = "Failed" ]; then
            status_color="\033[0;31m"
            failed_count=$((failed_count + 1))
        else
            passed_count=$((passed_count + 1))
        fi
        total_time=$((total_time + duration_ms))

        printf "  %-30s %b%-8b\033[0m  %s\n" "$name" "$status_color" "$status" "$duration_str"
    done
    echo "-----------------------------------------"
    local total_tests="${#TEST_NAMES[@]}"
    local total_time_str
    total_time_str=$(format_duration "$total_time")
    printf "  total: %d tests, %d passed, %d failed\n" "$total_tests" "$passed_count" "$failed_count"
    printf "  total execution time: %s\n" "$total_time_str"
    echo "========================================="
}

SUCCESS=0
rm -rf /tmp/test-spindle-microvm-logs

log "setup local cache & temp environment"
TEMP_DIR=$(mktemp -d -t test-spindle-microvm-XXXXXX)

log "build spindle & microvm image tarball"
nix develop --command go build -o spindle/spindle-microvm-run ./cmd/spindle-microvm-run
TARBALL_PATH=$(nix build .#spindle-nixos-image-tarball --no-link --print-out-paths)
mkdir -p "$TEMP_DIR/image"
tar -C "$TEMP_DIR/image" -xzf "$TARBALL_PATH"
IMAGE_SPEC_JSON="$TEMP_DIR/image/spindle-nixos-x86_64-linux.json"

log "build alpine microvm image tarball"
ALPINE_TARBALL_PATH=$(nix build .#spindle-alpine-image-tarball --no-link --print-out-paths)
mkdir -p "$TEMP_DIR/alpine-image"
tar -C "$TEMP_DIR/alpine-image" -xzf "$ALPINE_TARBALL_PATH"
ALPINE_IMAGE_SPEC_JSON="$TEMP_DIR/alpine-image/spindle-alpine-x86_64-linux.json"
cleanup() {
    print_summary
    log "cleaning up..."
    if [ -f "$TEMP_DIR/ncps.pid" ]; then
        kill "$(cat "$TEMP_DIR/ncps.pid")" || true
    fi
    if [ "$SUCCESS" -ne 1 ]; then
        echo "test failed. copying logs to /tmp/test-spindle-microvm-logs"
        mkdir -p /tmp/test-spindle-microvm-logs
        for f in "$TEMP_DIR"/*; do
            if [ -f "$f" ]; then
                if [[ "$f" == *.log ]]; then
                    strip_ansi "$f" > "/tmp/test-spindle-microvm-logs/$(basename "$f")"
                else
                    cp "$f" /tmp/test-spindle-microvm-logs/
                fi
            fi
        done
    fi
    chmod -R +w "$TEMP_DIR" || true
    rm -rf "$TEMP_DIR"
    echo "done"
}
trap cleanup EXIT

./spindle/engines/microvm/start-test-cache.sh "$TEMP_DIR" 8501
source "$TEMP_DIR/env.sh"

run_vm() {
    local name=""
    local timeout="60s"
    local upload=0
    local activate=""
    local no_cache=0
    local db=""
    local spec="$IMAGE_SPEC_JSON"

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --spec)
                spec="$2"
                shift 2
                ;;
            --name)
                name="$2"
                shift 2
                ;;
            --timeout)
                timeout="$2"
                shift 2
                ;;
            --upload)
                upload=1
                shift
                ;;
            --activate)
                activate="$2"
                shift 2
                ;;
            --no-cache)
                no_cache=1
                shift
                ;;
            --db)
                db="$2"
                shift 2
                ;;
            --)
                shift
                break
                ;;
            *)
                echo "unknown argument: $1" >&2
                exit 1
                ;;
        esac
    done

    local args=(
        --image-spec "$spec"
        --serial-log "$TEMP_DIR/serial-${name}.log"
        --qemu-log "$TEMP_DIR/qemu-${name}.log"
        --exec-timeout "$timeout"
    )

    if [ "$no_cache" -eq 0 ]; then
        args+=(
            --cache-read-url "$CACHE_URL"
            --cache-trusted-public-key "$CACHE_PUBKEY"
        )
    fi

    if [ "$upload" -eq 1 ]; then
        args+=(
            --cache-upload-url "$CACHE_UPLOAD_URL?secret-key=$CACHE_SECRET_KEY_PATH"
        )
    fi

    if [ -n "$activate" ]; then
        args+=(
            --activate-config "$activate"
        )
    fi

    if [ -n "$db" ]; then
        args+=(
            --db "$db"
        )
    fi

    local out
    if ! out=$(spindle/spindle-microvm-run "${args[@]}" -- "$@" 2>&1); then
        echo "$out" | strip_ansi >&2
        strip_ansi "$TEMP_DIR/serial-${name}.log" >&2
        strip_ansi "$TEMP_DIR/qemu-${name}.log" >&2
        return 1
    fi
    echo "$out"
}

test_case() {
    local name="$1"
    local func="$2"

    if [ -n "${TEST_ONLY:-}" ] && [ "${TEST_ONLY}" != "$name" ]; then
        return 0
    fi

    log "running test case: $name"
    local start_time
    start_time=$(get_time_ms)

    local status="Passed"
    if ! "$func"; then
        status="Failed"
    fi

    local end_time
    end_time=$(get_time_ms)
    local duration_ms=$((end_time - start_time))

    TEST_NAMES+=("$name")
    TEST_STATUSES+=("$status")
    TEST_TIMES+=("$duration_ms")

    local duration_str
    duration_str=$(format_duration "$duration_ms")

    if [ "$status" = "Failed" ]; then
        echo "error: test case $name failed (took $duration_str)"
        exit 1
    else
        log "test case $name passed (took $duration_str)"
    fi
}

test_realize() {
    local test_store_path
    test_store_path=$(nix-build -E 'with import <nixpkgs> {}; writeText "test-file" "hello from cache"' --no-out-link)
    nix copy --to "$CACHE_UPLOAD_URL?secret-key=$CACHE_SECRET_KEY_PATH" "$test_store_path"

    local out
    out=$(run_vm --name "realize" --timeout "60s" -- /run/current-system/sw/bin/bash -lc '
set -euo pipefail
store_path=$1
cache_url=$(sed -n "s/^extra-substituters = //p" /run/spindle/nix.conf)
cache_url=${cache_url%% *}
if [ -z "$cache_url" ]; then
    echo "error: cache URL not found in /run/spindle/nix.conf" >&2
    exit 1
fi

http_version=$(/run/current-system/sw/bin/curl --http2-prior-knowledge -fsS -o /dev/null -w "%{http_version}" "$cache_url/nix-cache-info")
echo "http_version=$http_version"
case "$http_version" in
    2|2.0) ;;
    *)
        echo "error: cache proxy did not negotiate HTTP/2 (got $http_version)" >&2
        exit 1
        ;;
esac

/run/current-system/sw/bin/nix-store --realise "$store_path" >/dev/null
' bash "$test_store_path") || return 1

    if ! echo "$out" | strip_ansi | grep -q -E "^http_version=2(\\.0)?$"; then
        echo "error: cache proxy did not report HTTP/2" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi
    echo "success: store path realized from cache and cache proxy accepted cleartext HTTP/2"
}

test_build_upload() {
    local nix_expr='with import <nixpkgs> {}; writeText "uploaded-test-file" "hello from vm upload"'
    local out
    out=$(run_vm --name "build-upload" --timeout "120s" --upload -- /run/current-system/sw/bin/bash -l -c "nix-build -E '$nix_expr' --no-out-link") || return 1

    local built_path
    built_path=$(echo "$out" | strip_ansi | grep -v '\.drv' | grep -o '/nix/store/[a-z0-9]*-uploaded-test-file' | head -n 1 || true)
    if [ -z "$built_path" ]; then
        echo "error: could not find built store path in vm output" >&2
        return 1
    fi
    echo "extracted path: $built_path"

    local hash
    hash=$(basename "$built_path" | cut -d'-' -f1)
    if ! curl -s -f "http://127.0.0.1:8501/${hash}.narinfo" > /dev/null; then
        echo "error: built store path was not uploaded to the binary cache" >&2
        return 1
    fi
    echo "success: store path uploaded to cache"
}

test_networking() {
    local hello_path
    hello_path=$(nix-build -E 'with import <nixpkgs> {}; hello' --no-out-link)

    local out
    out=$(run_vm --name "networking" --timeout "120s" --no-cache -- /run/current-system/sw/bin/bash -c "/run/current-system/sw/bin/curl -I --connect-timeout 1 -m 1 http://10.0.2.2:8501; /run/current-system/sw/bin/nix-store --realise $hello_path") || return 1

    if echo "$out" | grep -qi -E "unreachable|timeout|failed to connect|timed out" || echo "$out" | grep -q "exited with code"; then
        echo "success: host network access blocked"
    else
        echo "error: guest vm accessed host network or returned unexpected output" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi
    echo "success: guest vm reached the internet and substituted hello"
}

test_nixpkgs_substitution_and_no_upload() {
    local hello_path
    hello_path=$(nix-build -E 'with import <nixpkgs> {}; hello' --no-out-link)

    local out
    out=$(run_vm --name "nixpkgs-hello" --timeout "180s" --upload -- /run/current-system/sw/bin/nix-store --realise "$hello_path") || return 1

    # Check that it was substituted from our proxy
    if ! echo "$out" | strip_ansi | grep -q -E "copying path.*hello"; then
        echo "error: hello package was not substituted (or output mismatch)" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    # Check that nothing was uploaded to the cache
    if ! echo "$out" | strip_ansi | grep -q "cache uploaded: 0"; then
        echo "error: hello package substitution triggered cache upload" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    echo "success: hello package substituted from upstream cache and was not uploaded"
}

test_activation() {
    local user_config='{"services":{"openssh":{"enable":true}}}'
    local db_path="$TEMP_DIR/spindle.db"
    local out

    local t1_start
    t1_start=$(get_time_ms)

    # compile the config and saves it to the cache db
    out=$(run_vm --name "activation-first" --timeout "600s" --activate "$user_config" --db "$db_path" --upload -- /run/current-system/sw/bin/systemctl is-active sshd) || return 1

    local t1_end
    t1_end=$(get_time_ms)
    local t1_dur=$((t1_end - t1_start))
    local t1_str
    t1_str=$(format_duration "$t1_dur")
    echo "info: activation-first took $t1_str"

    if ! echo "$out" | strip_ansi | grep -q "^active$"; then
        echo "error: sshd not active after first config activation" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    local t2_start
    t2_start=$(get_time_ms)

    # should pull the built configuration from the cache
    out=$(run_vm --name "activation-second" --timeout "600s" --activate "$user_config" --db "$db_path" -- /run/current-system/sw/bin/systemctl is-active sshd) || return 1

    local t2_end
    t2_end=$(get_time_ms)
    local t2_dur=$((t2_end - t2_start))
    local t2_str
    t2_str=$(format_duration "$t2_dur")
    echo "info: activation-second took $t2_str"

    if ! echo "$out" | strip_ansi | grep -q "^active$"; then
        echo "error: sshd not active after cached config activation" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    if ! echo "$out" | strip_ansi | grep -q "realizing cached NixOS config"; then
        echo "error: second run did not realize cached configuration" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi
    echo "success: user config activated and successfully realized from cache (first run: $t1_str, second run: $t2_str)"
}

test_dependencies_and_services() {
    local user_config='{
        "registry": {
            "my-nixpkgs": "nixpkgs"
        },
        "dependencies": [
            "my-nixpkgs#hello"
        ],
        "services": {
            "openssh": {
                "enable": true,
                "authorizedKeysFiles": [
                    "/etc/ssh/authorized_keys"
                ]
            },
            "udev": {
                "enable": true,
                "packages": ["my-nixpkgs#hello"]
            }
        },
        "virtualisation": {
            "docker": {
                "enable": true
            }
        }
    }'
    local out

    out=$(run_vm --name "deps-and-services" --timeout "600s" --activate "$user_config" -- /run/current-system/sw/bin/bash -l -c "systemctl is-active sshd && systemctl is-active docker && hello") || return 1

    if ! echo "$out" | strip_ansi | grep -q "Hello, world!"; then
        echo "error: hello command not found or failed to execute" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi
    if ! echo "$out" | strip_ansi | grep -q "^active$"; then
        echo "error: sshd or docker not active" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi
    echo "success: top-level dependencies, registry config, package list conversion, and implicit services verified successfully"
}

test_alpine() {
    local out
    out=$(run_vm --spec "$ALPINE_IMAGE_SPEC_JSON" --name "alpine" --timeout "120s" --no-cache -- /bin/sh -lc '
set -eu
echo "release=$(cat /etc/alpine-release)"
echo "user=$(id -un)"
git version
bash -c "echo bash=\$BASH_VERSION"
touch /workspace/write-test
echo "workspace writable"
git ls-remote https://tangled.org/@tangled.org/core HEAD >/dev/null
echo "git over https ok"
apk add make
echo "apk ok"
') || return 1

    echo "$out" | strip_ansi >&2
    for needle in "release=" "user=spindle-workflow" "git version" "bash=" "workspace writable" "git over https ok" "apk ok"; do
        if ! echo "$out" | strip_ansi | grep -q "$needle"; then
            echo "error: alpine guest output missing $needle" >&2
            return 1
        fi
    done
    echo "success: alpine guest booted, ran as workflow user, wrote workspace, cloned over https"
}

test_alpine_nix() {
    local test_store_path
    test_store_path=$(nix-build -E 'with import <nixpkgs> {}; writeText "alpine-nix-test" "hello from cache to alpine"' --no-out-link)
    nix copy --to "$CACHE_UPLOAD_URL?secret-key=$CACHE_SECRET_KEY_PATH" "$test_store_path"

    local out
    out=$(run_vm --spec "$ALPINE_IMAGE_SPEC_JSON" --name "alpine-nix" --timeout "180s" --upload -- /bin/sh -lc '
set -eu
export HOME=/workspace
store_path=$1
nix-store --realise "$store_path" >/dev/null
echo "substituted=$(cat "$store_path")"
built=$(nix-build -E "derivation { name = \"alpine-built\"; system = \"x86_64-linux\"; builder = \"/bin/sh\"; args = [ \"-c\" \"echo built-on-alpine > \$out\" ]; }" --no-out-link)
echo "built=$built"
' sh "$test_store_path") || return 1

    if ! echo "$out" | strip_ansi | grep -q "substituted=hello from cache to alpine"; then
        echo "error: alpine guest did not substitute the test path from the cache" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    local built_path
    built_path=$(echo "$out" | strip_ansi | grep -o 'built=/nix/store/[a-z0-9]*-alpine-built' | cut -d= -f2)
    if [ -z "$built_path" ]; then
        echo "error: could not find built store path in alpine guest output" >&2
        echo "$out" | strip_ansi >&2
        return 1
    fi

    local hash
    hash=$(basename "$built_path" | cut -d'-' -f1)
    if ! curl -s -f "http://127.0.0.1:8501/${hash}.narinfo" > /dev/null; then
        echo "error: alpine-built store path was not uploaded to the binary cache" >&2
        return 1
    fi
    echo "success: alpine guest substituted from cache, built a derivation, and uploaded it"
}

test_case "alpine" test_alpine
test_case "alpine-nix" test_alpine_nix
test_case "realize" test_realize
test_case "build-upload" test_build_upload
test_case "networking" test_networking
test_case "substitution-no-upload" test_nixpkgs_substitution_and_no_upload
test_case "activation" test_activation
test_case "deps-and-services" test_dependencies_and_services

SUCCESS=1
log "passed!!"
