#!/bin/bash
# ══════════════════════════════════════════════════════════════════════════════
# Volt Hybrid Integration Test Helpers
#
# Shared functions for all hybrid integration test scripts.
# Source this file at the top of every test:
#   source "$(dirname "$0")/test_helpers.sh"
# ══════════════════════════════════════════════════════════════════════════════

set -uo pipefail

# ── Configuration ─────────────────────────────────────────────────────────────

TEST_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VOLT="${VOLT:-$(cd "$TEST_DIR/../.." && pwd)/volt}"
MANIFEST_DIR="$TEST_DIR/test-manifests"
# Prefix all test workload names so cleanup can nuke them
TEST_PREFIX="volt-inttest"
# Timeout for workload operations (seconds)
OP_TIMEOUT="${OP_TIMEOUT:-60}"
# Timeout for workload boot readiness (seconds)
BOOT_TIMEOUT="${BOOT_TIMEOUT:-30}"

# ── Counters ──────────────────────────────────────────────────────────────────

PASS=0
FAIL=0
SKIP=0
TOTAL=0
ERRORS=""
CLEANUP_WORKLOADS=()

# ── Color / Formatting ───────────────────────────────────────────────────────

if [[ -t 1 ]]; then
    GREEN='\033[0;32m'
    RED='\033[0;31m'
    YELLOW='\033[0;33m'
    DIM='\033[0;90m'
    BOLD='\033[1m'
    RESET='\033[0m'
else
    GREEN='' RED='' YELLOW='' DIM='' BOLD='' RESET=''
fi

# ── Test Primitives ──────────────────────────────────────────────────────────

# Pass a test with a description
pass() {
    local desc="$1"
    PASS=$((PASS + 1))
    TOTAL=$((TOTAL + 1))
    echo -e "  ${GREEN}✓${RESET} $desc"
}

# Fail a test with a description and optional detail
fail() {
    local desc="$1"
    local detail="${2:-}"
    FAIL=$((FAIL + 1))
    TOTAL=$((TOTAL + 1))
    echo -e "  ${RED}✗${RESET} $desc"
    if [[ -n "$detail" ]]; then
        echo -e "    ${DIM}→ $detail${RESET}"
    fi
    ERRORS="${ERRORS}\n  ✗ $desc${detail:+: $detail}"
}

# Skip a test
skip() {
    local desc="$1"
    local reason="${2:-}"
    SKIP=$((SKIP + 1))
    TOTAL=$((TOTAL + 1))
    echo -e "  ${YELLOW}⊘${RESET} $desc (skipped${reason:+: $reason})"
}

# Assert a command succeeds (exit 0)
assert_ok() {
    local desc="$1"
    shift
    local output
    if output=$("$@" 2>&1); then
        pass "$desc"
        return 0
    else
        fail "$desc" "exit=$?, output: $(echo "$output" | head -3 | tr '\n' ' ')"
        return 1
    fi
}

# Assert a command fails (non-zero exit)
assert_fail() {
    local desc="$1"
    shift
    local output
    if output=$("$@" 2>&1); then
        fail "$desc" "expected failure but got exit=0"
        return 1
    else
        pass "$desc (correctly fails)"
        return 0
    fi
}

# Assert command output contains a string (case-insensitive)
assert_contains() {
    local desc="$1"
    local expected="$2"
    shift 2
    local output
    if output=$("$@" 2>&1) && echo "$output" | grep -qi "$expected"; then
        pass "$desc"
        return 0
    else
        fail "$desc" "expected output to contain '$expected'"
        return 1
    fi
}

# Assert command output does NOT contain a string
assert_not_contains() {
    local desc="$1"
    local unexpected="$2"
    shift 2
    local output
    output=$("$@" 2>&1) || true
    if echo "$output" | grep -qi "$unexpected"; then
        fail "$desc" "output should not contain '$unexpected'"
        return 1
    else
        pass "$desc"
        return 0
    fi
}

# Assert two values are equal
assert_eq() {
    local desc="$1"
    local expected="$2"
    local actual="$3"
    if [[ "$expected" == "$actual" ]]; then
        pass "$desc"
        return 0
    else
        fail "$desc" "expected='$expected', actual='$actual'"
        return 1
    fi
}

# Assert a value is non-empty
assert_nonempty() {
    local desc="$1"
    local value="$2"
    if [[ -n "$value" ]]; then
        pass "$desc"
        return 0
    else
        fail "$desc" "expected non-empty value"
        return 1
    fi
}

# Assert a file exists
assert_file_exists() {
    local desc="$1"
    local path="$2"
    if [[ -f "$path" ]]; then
        pass "$desc"
        return 0
    else
        fail "$desc" "file not found: $path"
        return 1
    fi
}

# Assert a directory exists
assert_dir_exists() {
    local desc="$1"
    local path="$2"
    if [[ -d "$path" ]]; then
        pass "$desc"
        return 0
    else
        fail "$desc" "directory not found: $path"
        return 1
    fi
}

# Assert a file does NOT exist
assert_no_file() {
    local desc="$1"
    local path="$2"
    if [[ ! -e "$path" ]]; then
        pass "$desc"
        return 0
    else
        fail "$desc" "expected $path to not exist"
        return 1
    fi
}

# ── Workload Helpers ─────────────────────────────────────────────────────────

# Generate a unique workload name with the test prefix
test_name() {
    local base="$1"
    echo "${TEST_PREFIX}-${base}-$$"
}

# Register a workload for cleanup on exit
register_cleanup() {
    local name="$1"
    CLEANUP_WORKLOADS+=("$name")
}

# Create a container workload from image — returns immediately
create_container() {
    local name="$1"
    local image="${2:-/var/lib/volt/images/ubuntu_24.04}"
    local extra_flags="${3:-}"
    register_cleanup "$name"
    # shellcheck disable=SC2086
    sudo "$VOLT" container create --name "$name" --image "$image" --backend hybrid $extra_flags 2>&1
}

# Start a workload and wait until it's running
start_workload() {
    local name="$1"
    sudo "$VOLT" container start "$name" 2>&1
}

# Stop a workload
stop_workload() {
    local name="$1"
    sudo "$VOLT" container stop "$name" 2>&1
}

# Destroy a workload (stop + delete)
destroy_workload() {
    local name="$1"
    sudo "$VOLT" container delete "$name" --force 2>&1 || true
}

# Execute a command inside a running container
exec_in() {
    local name="$1"
    shift
    sudo "$VOLT" container exec "$name" -- "$@" 2>&1
}

# Wait for a container to be "running" according to machinectl/systemd
wait_running() {
    local name="$1"
    local timeout="${2:-$BOOT_TIMEOUT}"
    local elapsed=0
    while (( elapsed < timeout )); do
        if sudo machinectl show "$name" --property=State 2>/dev/null | grep -q "running"; then
            return 0
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done
    return 1
}

# Wait for systemd inside a boot-mode container to reach a target
wait_booted() {
    local name="$1"
    local timeout="${2:-$BOOT_TIMEOUT}"
    local elapsed=0
    while (( elapsed < timeout )); do
        if sudo machinectl shell "$name" /bin/systemctl is-system-running 2>/dev/null | grep -qE "running|degraded"; then
            return 0
        fi
        sleep 1
        elapsed=$((elapsed + 1))
    done
    return 1
}

# Get the leader PID of a running container
get_leader_pid() {
    local name="$1"
    sudo machinectl show "$name" --property=Leader --value 2>/dev/null | tr -d '[:space:]'
}

# Get the IP address of a running container
get_container_ip() {
    local name="$1"
    sudo machinectl show "$name" --property=Addresses --value 2>/dev/null | awk '{print $1}'
}

# Check if a container rootfs directory exists
rootfs_exists() {
    local name="$1"
    [[ -d "/var/lib/volt/containers/$name" ]] || [[ -d "/var/lib/machines/$name" ]]
}

# Get the systemd unit name for a hybrid container
hybrid_unit() {
    local name="$1"
    echo "volt-hybrid@${name}.service"
}

# ── Prerequisite Checks ─────────────────────────────────────────────────────

# Check if running as root (required for nspawn operations)
require_root() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: These integration tests require root (systemd-nspawn needs it)."
        echo "Run with: sudo ./run_tests.sh"
        exit 1
    fi
}

# Check if a base image is available
require_image() {
    local image_path="${1:-/var/lib/volt/images/ubuntu_24.04}"
    if [[ ! -d "$image_path" ]]; then
        echo "ERROR: Base image not found at $image_path"
        echo "Create one with: sudo debootstrap noble $image_path http://archive.ubuntu.com/ubuntu"
        return 1
    fi
    return 0
}

# Check if systemd-nspawn is available
require_nspawn() {
    if ! command -v systemd-nspawn &>/dev/null; then
        echo "ERROR: systemd-nspawn not found. Install systemd-container."
        return 1
    fi
    return 0
}

# Check if volt binary exists and is executable
require_volt() {
    if [[ ! -x "$VOLT" ]]; then
        echo "ERROR: volt binary not found or not executable at $VOLT"
        return 1
    fi
    return 0
}

# ── Cleanup ──────────────────────────────────────────────────────────────────

# Clean up all registered test workloads
cleanup_all() {
    local exit_code=$?
    echo ""
    echo -e "${DIM}Cleaning up test workloads...${RESET}"
    for name in "${CLEANUP_WORKLOADS[@]}"; do
        if sudo machinectl show "$name" &>/dev/null 2>&1; then
            sudo machinectl terminate "$name" &>/dev/null 2>&1 || true
            sleep 1
        fi
        sudo systemctl stop "volt-hybrid@${name}.service" &>/dev/null 2>&1 || true
        sudo systemctl stop "systemd-nspawn@${name}.service" &>/dev/null 2>&1 || true
        # Remove rootfs
        sudo rm -rf "/var/lib/volt/containers/$name" 2>/dev/null || true
        sudo rm -rf "/var/lib/machines/$name" 2>/dev/null || true
        # Remove unit files
        sudo rm -f "/etc/systemd/system/volt-hybrid@${name}.service" 2>/dev/null || true
        sudo rm -f "/etc/systemd/nspawn/${name}.nspawn" 2>/dev/null || true
    done
    sudo systemctl daemon-reload &>/dev/null 2>&1 || true
    echo -e "${DIM}Cleanup complete.${RESET}"
    return $exit_code
}

# ── Results Summary ──────────────────────────────────────────────────────────

print_results() {
    local suite_name="${1:-Hybrid Integration Tests}"
    echo ""
    echo "════════════════════════════════════════════════════════════════"
    echo -e "${BOLD}$suite_name${RESET}"
    echo "────────────────────────────────────────────────────────────────"
    echo -e "  Passed:  ${GREEN}${PASS}${RESET}"
    echo -e "  Failed:  ${RED}${FAIL}${RESET}"
    if [[ $SKIP -gt 0 ]]; then
        echo -e "  Skipped: ${YELLOW}${SKIP}${RESET}"
    fi
    echo "  Total:   ${TOTAL}"
    echo "════════════════════════════════════════════════════════════════"

    if [[ $FAIL -gt 0 ]]; then
        echo ""
        echo -e "${RED}Failures:${RESET}"
        echo -e "$ERRORS"
        return 1
    fi

    echo ""
    echo -e "${GREEN}All tests passed! ✅${RESET}"
    return 0
}

# ── Section Header ───────────────────────────────────────────────────────────

section() {
    local title="$1"
    echo ""
    echo -e "${BOLD}${title}${RESET}"
}