#!/usr/bin/env bash
set -euo pipefail

# Path to the folder containing this script.
script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Get current working directory
cwd="$(pwd)"

# If script_dir is the same as cwd, set script_rel_path to "./$(basename "$0")"
if [[ "$script_dir" == "$cwd" ]]; then
    script_rel_path="./$(basename "$0")"
else
    script_rel_path="${script_dir#"$cwd"/}/$(basename "$0")"
fi

usage() {
    cat <<EOF
Usage: ${script_rel_path} [COMMAND] [OPTIONS] [ARGS]

Options:
  -h, --help             Display this help message and exit.
  -c, --concurrency <N>  Limit concurrency to N parallel jobs. If not specified,
                         attempts to automatically detect and use a sane maximum 
                         based on both CPU cores and total RAM.
  -u, --update           Update snapshots. Use this to update snapshots after
                         making changes to the code that affect the tests.

Commands:
  diff                   Start the snapshot diff tool. This tool allows you to
                         compare the current snapshots with the expected
                         snapshots for the previous test run. This is useful for
                         identifying changes in the output that may be
                         intentional or not.

Arguments:
  LOCALES                Specify one or more locales to test. If none are
                         provided, all locales will be tested.

Examples:
  ${script_rel_path} pl nn          # Test Polish and Nynorsk
  ${script_rel_path} -c 2 pl nn fr  # Test 3 locales, maximum of 2 parallel jobs
  ${script_rel_path} -u             # Update snapshots for all locales
  ${script_rel_path} diff           # Start the snapshot diff tool
  ${script_rel_path} --help         # Display this help message
  ${script_rel_path} diff --help    # Display help for the diff tool
EOF
}

# Return the number of available CPU cores (cross‐platform).
detect_cpu_cores() {
    # Linux + WSL
    if command -v nproc >/dev/null 2>&1; then
        nproc
    # macOS
    elif [[ "$(uname -s)" == "Darwin" ]]; then
        # 'sysctl -n hw.logicalcpu' is usually the right approach on modern macOS
        sysctl -n hw.logicalcpu 2>/dev/null || echo 2
    else
        # Fall back to a 'safe' default if we can't detect
        echo 2
    fi
}

# Try to determine the total system memory in GB (integer). If detection fails,
# fall back to 4 GB as a safe guess.
detect_memory_in_gb() {
    # Linux: /proc/meminfo
    if [[ -r /proc/meminfo ]]; then
        local mem_kb
        mem_kb="$(awk '/MemTotal:/ {print $2}' /proc/meminfo)"
        # Convert kB -> GB. 1 GB = 1,048,576 kB
        echo $(( mem_kb / 1048576 ))
        return
    fi

    # macOS: sysctl hw.memsize
    if command -v sysctl >/dev/null 2>&1; then
        local mem_bytes
        mem_bytes="$(sysctl -n hw.memsize 2>/dev/null || echo 0)"
        # Convert bytes -> GB. 1 GB = 1,073,741,824 bytes
        echo $(( mem_bytes / 1073741824 ))
        return
    fi

    # If memory cannot be detected at all, fall back to a reasonable default.
    echo 4
}

# Combine CPU core count + total RAM to produce a default concurrency value.
# Heuristic:
#   concurrency_from_ram = max(2, floor(RAM_in_GB / 16))
#   concurrency_final    = min(concurrency_from_ram, cpu_cores)
detect_auto_concurrency() {
    local cores mem_gb concurrency_ram concurrency_final

    cores="$(detect_cpu_cores)"
    mem_gb="$(detect_memory_in_gb)"

    # Rule of thumb: 1 heavy build/test per 16GB, but never less than 2.
    concurrency_ram=$(( mem_gb / 16 ))
    if [[ $concurrency_ram -lt 2 ]]; then
        concurrency_ram=2
    fi

    # Final concurrency is limited by the number of CPU cores.
    if [[ "$cores" -lt "$concurrency_ram" ]]; then
        concurrency_final="$cores"
    else
        concurrency_final="$concurrency_ram"
    fi

    echo "$concurrency_final"
}

concurrency=""
locales=()
update_snapshots=false

# Parse command-line arguments.
if [[ "${1:-}" == "diff" ]]; then
    # Shift off "diff", so any subsequent arguments can be passed as well
    shift
    # Call an external script that handles the TUI logic.
    if [ ! -x "${script_dir}/smoketest-diff.sh" ]; then
        echo "Error: smoketest-diff.sh not found or not executable." >&2
        exit 1
    fi
    exec "${script_dir}/smoketest-diff.sh" "$@"
fi

while [[ "$#" -gt 0 ]]; do
    case "$1" in
        -h|--help)
            usage
            exit 0
            ;;
        -c|--concurrency)
            shift
            concurrency="${1:-}"
            if [[ -z "$concurrency" || "$concurrency" =~ ^- ]]; then
                echo "Error: --concurrency requires a numeric argument." >&2
                exit 1
            fi
            shift
            ;;
        -u|--update)
            update_snapshots=true
            shift
            ;;
        -*)
            echo "Unknown option: $1" >&2
            usage
            exit 1
            ;;
        *)
            locales+=("$1")
            shift
            ;;
    esac
done

# If concurrency wasn't provided, attempt to auto-detect.
if [[ -z "$concurrency" ]]; then
    concurrency="$(detect_auto_concurrency)"
fi

# Validate concurrency is a positive integer.
if ! [[ "$concurrency" =~ ^[0-9]+$ ]]; then
    echo "Error: concurrency ('$concurrency') is not a valid integer." >&2
    exit 1
fi
if [[ "$concurrency" -lt 1 ]]; then
    echo "Error: concurrency must be >= 1." >&2
    exit 1
fi

# If no locales were specified, use all locales.
if [ ${#locales[@]} -eq 0 ]; then
    # Get list of all folders in the locale folder that don't start with _
    locales=($(find "${script_dir}/../locale" -mindepth 1 -maxdepth 1 -type d -exec basename {} \; | grep -v '^_'))
fi

# Verify that SSH_AUTH_SOCK is set.
if [ -z "${SSH_AUTH_SOCK:-}" ]; then
    echo "SSH_AUTH_SOCK is not set. Make sure you have an SSH key added to your SSH agent." >&2
    echo "You can start an SSH agent like so:" >&2
    echo >&2
    echo "  eval \$(ssh-agent)" >&2
    echo >&2
    echo "And then add your key with:" >&2
    echo >&2
    echo "  ssh-add /path/to/your/key" >&2
    echo >&2
    # Detect first available SSH key. If running on WSL, include $WIN_HOME/.ssh
    sshKeyPath="$(find "$HOME"/.ssh -type f -name 'id_*' ! -name '*.pub' 2>/dev/null | head -n 1)"
    if [ -z "$sshKeyPath" ] && [ -n "${WIN_HOME:-}" ]; then
        sshKeyPath=$(find "$WIN_HOME"/.ssh -type f -name 'id_*' ! -name '*.pub' 2>/dev/null | head -n 1)
    fi

    if [ -n "$sshKeyPath" ]; then
        echo "Detected SSH key at $sshKeyPath." >&2
        echo "If this is the key you want to use, add it to your SSH agent with:" >&2
        echo >&2
        echo "  ssh-add $sshKeyPath" >&2
        echo >&2
    fi
    exit 1
fi

# Verify that Docker is installed and running.
if ! command -v docker >/dev/null 2>&1; then
    echo "Docker is not installed. Please install Docker to run smoke tests." >&2
    exit 1
fi

if ! docker info >/dev/null 2>&1; then
    echo "Docker is not running. Please start Docker to run smoke tests." >&2
    exit 1
fi

imageName="pronounspage-smoke-tests"

# Make sure BuildKit is enabled (for SSH forwarding).
export DOCKER_BUILDKIT=1

# Grab USER_ID and GROUP_ID to pass into the build, if available.
if command -v id >/dev/null 2>&1; then
    USER_ID=$(id -u)
    GROUP_ID=$(id -g)
else
    USER_ID=""
    GROUP_ID=""
fi

echo "Building Docker image '$imageName' with SSH Agent Forwarding..." >&2
docker build \
    -t "$imageName" \
    -f "${script_dir}/../docker/smoketest.Dockerfile" \
    --ssh default \
    --build-arg USER_ID="$USER_ID" \
    --build-arg GROUP_ID="$GROUP_ID" \
    .

# Ensure the snapshots directory exists.
mkdir -p "${script_dir}/__snapshots__"

echo "Running tests for locales: ${locales[*]}" >&2
echo "Max containers running in parallel: $concurrency" >&2

LOG_DIR="$(mktemp -d -t pronounspage-logs-XXXXXX)"
FAIL_FILE="$LOG_DIR/.failures"
touch "$FAIL_FILE"

# Track status for each locale:
#   "pending" -> not started yet.
#   "running" -> docker container is running.
#   "passed"  -> test completed with exit 0.
#   "failed"  -> test completed with non‐zero exit.
declare -a pids=()
declare -a statuses=()

# Initialize all statuses to "pending"
for ((i=0; i<${#locales[@]}; i++)); do
    pids[i]=""
    statuses[i]="pending"
done

# For color-coding statuses in the "Running tests for locales" line.
COLOR_RED="\033[1;31m"
COLOR_GREEN="\033[1;32m"
COLOR_YELLOW="\033[1;33m"
COLOR_GRAY="\033[1;30m"
COLOR_RESET="\033[0m"

# Helper function to start a docker test run in background.
start_test() {
    local i="$1"
    local loc="${locales[$i]}"
    local log_file="$LOG_DIR/$loc.log"
    local update_flag=""

    statuses[i]="running"

    if [ "$update_snapshots" = true ]; then
        update_flag="-u"
    fi

    docker run --rm \
        -v "${script_dir}/__snapshots__:/app/test/__snapshots__" \
        -e TEST_LOCALES="$loc" \
        "$imageName" \
        $update_flag \
        >"$log_file" 2>&1 &

    pids[i]=$!
}

# Check if a test has completed and update its status accordingly.
check_test() {
    local i="$1"
    local pid="${pids[$i]}"

    # If it's not running, nothing to check.
    if [ "${statuses[$i]}" != "running" ]; then
        return
    fi

    # If kill -0 succeeds, process is still running.
    if kill -0 "$pid" 2>/dev/null; then
        return
    fi

    # Otherwise, it's finished. Check exit status for pass/fail.
    if ! wait "$pid"; then
        echo "${locales[$i]}" >> "$FAIL_FILE"
        statuses[i]="failed"
    else
        statuses[i]="passed"
    fi
}

print_elapsed_time() {
    local end_time
    end_time=$(date +%s)
    local elapsed_time
    elapsed_time=$((end_time - start_time))
    local hours
    hours=$((elapsed_time / 3600))
    local minutes
    minutes=$(( (elapsed_time % 3600) / 60 ))
    local seconds
    seconds=$((elapsed_time % 60))

    if [ "$hours" -gt 0 ]; then
        echo "Total time elapsed: ${hours}h ${minutes}m ${seconds}s" >&2
    elif [ "$minutes" -gt 0 ]; then
        echo "Total time elapsed: ${minutes}m ${seconds}s" >&2
    else
        echo "Total time elapsed: ${seconds}s" >&2
    fi
}

# Remove the logs directory only if all tests pass.
cancelled=false
cleanup() {
    # Restore cursor (i.e. make it visible again).
    tput cnorm

    # If run was cancelled, do not print success.
    if [ "$cancelled" = true ]; then
        echo >&2
        echo "Test run was cancelled; leaving $LOG_DIR in place for inspection." >&2
        print_elapsed_time
        exit 1
    fi

    # If any failures were recorded, keep logs.
    if [ -s "$FAIL_FILE" ]; then
        echo >&2
        echo "Some tests failed; leaving $LOG_DIR in place for inspection." >&2
        echo >&2
        while read -r loc; do
            echo "  - $loc (logs in $LOG_DIR/$loc.log)" >&2
        done < "$FAIL_FILE"
        echo >&2
        print_elapsed_time
        echo >&2
        echo "You can also run the diff tool to compare snapshots:" >&2
        echo "  ${script_rel_path} diff" >&2
        exit 1
    else
        rm -rf "$LOG_DIR"
        echo -e "\033[1;32mAll locales passed successfully!\033[0m" >&2
        print_elapsed_time
    fi
    exit 0
}

# If Ctrl+C, kill all running containers, skip removing logs, restore cursor.
handle_sigint() {
    echo >&2
    echo "Cancelling all tests..." >&2
    # Kill all Docker runs.
    for ((i=0; i<${#locales[@]}; i++)); do
        if [ "${statuses[$i]}" = "running" ]; then
            kill "${pids[$i]}" 2>/dev/null || true
        fi
    done

    cancelled=true
    # Restore cursor.
    tput cnorm
    echo "Exiting..." >&2
    exit 1
}

trap cleanup EXIT
trap handle_sigint INT

# Hide cursor to make the interface look nicer.
tput civis

# Clears the screen and prints a summary line of all locales with color-coded
# statuses (pending in grey, passed in green, failed in red, running in yellow).
# Then, only prints the last few lines of logs for locales that are currently
# running, to avoid displaying overwhelming output.
update_tui() {
    clear
    tput cup 0 0

    # Print color-coded statuses for each locale in a single summary line.
    echo -n "Running tests for locales:" >&2
    for ((k=0; k<${#locales[@]}; k++)); do
        local loc="${locales[$k]}"
        case "${statuses[$k]}" in
            pending)
                echo -en " ${COLOR_GRAY}${loc}${COLOR_RESET}" >&2
                ;;
            running)
                echo -en " ${COLOR_YELLOW}${loc}${COLOR_RESET}" >&2
                ;;
            passed)
                echo -en " ${COLOR_GREEN}${loc}${COLOR_RESET}" >&2
                ;;
            failed)
                echo -en " ${COLOR_RED}${loc}${COLOR_RESET}" >&2
                ;;
        esac
    done
    echo >&2

    echo "Max containers running in parallel: $concurrency" >&2
    if [ "$update_snapshots" = true ]; then
        echo -e "${COLOR_YELLOW}Updating snapshots.${COLOR_RESET}" >&2
    fi
    echo >&2
    echo -e "\033[1;35mRunning test containers in parallel...\033[0m" >&2
    echo >&2

    # Print only currently running locales with last few lines of log.

    # Determine a sane number of lines to display per locale. Determine based on
    # the number of concurrent tests and the terminal height.

    # Each locale takes at least 3 lines on top of the log output. Additionally,
    # the header takes 2-3 lines. We want to leave some space for the header and
    # footer, so we'll subtract 3 from the terminal height to get the max lines,
    # and then calculate the remaining room for log output by dividing by the
    # number of concurrent tests and subtracting 3 again. Then we take the max
    # of that and 2 to ensure we always show at least one line (since the last
    # line may be blank).
    
    # Formula:
    #   max_lines = max(2, (term_height - 3) / concurrency - 3)

    local term_height
    term_height=$(tput lines)
    local max_lines
    max_lines=$(( (term_height - 3) / concurrency - 3 ))
    if [ "$max_lines" -lt 2 ]; then
        max_lines=2
    fi

    for ((k=0; k<${#locales[@]}; k++)); do
        if [ "${statuses[$k]}" = "running" ]; then
            local loc="${locales[$k]}"
            local log_file="$LOG_DIR/$loc.log"

            echo "==> $loc: Running..." >&2
            if [ -f "$log_file" ]; then
                tail -n "$max_lines" "$log_file" | while IFS= read -r line; do
                    local term_width
                    term_width=$(tput cols)
                    if [ ${#line} -gt $((term_width-5)) ]; then
                        echo "    ${line:0:$((term_width-5))}…" >&2
                    else
                        echo "    $line" >&2
                    fi
                done
            fi
            echo >&2
        fi
    done
}

# Main concurrency loop.
total="${#locales[@]}"
completed_count=0
running_count=0
i=0
start_time=$(date +%s)

while [ "$completed_count" -lt "$total" ]; do
    # Start new tests if we have capacity.
    while [ "$i" -lt "$total" ] && [ "$running_count" -lt "$concurrency" ]; do
        start_test "$i"
        running_count=$((running_count + 1))
        i=$((i + 1))
    done

    # Check for completed tests.
    for ((j=0; j<${#locales[@]}; j++)); do
        if [ "${statuses[$j]}" = "running" ]; then
            check_test "$j"
            if [ "${statuses[$j]}" = "passed" ] || [ "${statuses[$j]}" = "failed" ]; then
                running_count=$((running_count - 1))
                completed_count=$((completed_count + 1))
            fi
        fi
    done

    # Update the interface.
    update_tui

    # If we're not done, wait a bit before re-checking.
    if [ "$completed_count" -lt "$total" ]; then
        sleep 1
    fi
done

# One final update for good measure.
update_tui