#!/usr/bin/env bash

#########################################################################################
#### k6-test-runner.sh — Run k6 load tests with Prometheus push and formatted        ####
#### reports. Execute test scripts, push metrics, compare runs, threshold checks     ####
#### Requires: bash 4+, k6                                                           ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./k6-test-runner.sh --run ./tests/load.js                                     ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Colors (pre-initialized) ─────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""

setup_colors() {
    if [[ "${COLOR:-auto}" == "never" ]]; then
        return
    fi
    if [[ "${COLOR:-auto}" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }
die()     { err "$*"; exit 1; }

section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

field() {
    printf "  ${BOLD}%-22s${RESET} %s\n" "$1" "$2"
}

field_color() {
    printf "  ${BOLD}%-22s${RESET} %b\n" "$1" "$2"
}

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

# ── Defaults ──────────────────────────────────────────────────────────
RUN_MODE=""
TEST_SCRIPT=""
K6_PATH="${K6_PATH:-k6}"
VUS="${K6_VUS:-10}"
DURATION="${K6_DURATION:-30s}"
PUSH_GW="${K6_PUSH_GATEWAY:-}"
RESULTS_DIR="${K6_RESULTS_DIR:-./k6-results}"
OUTPUT_FORMAT="${K6_FORMAT:-text}"
THRESHOLDS_FILE=""
COMPARE_RUN=""
LIST_DIR=""
INSPECT_SCRIPT=""
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"
declare -a ENV_VARS=()
TAGS="${K6_TAGS:-}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""
RUN_ID=""

# ── Dependency checks ────────────────────────────────────────────────
require_k6() {
    if ! command -v "$K6_PATH" &>/dev/null; then
        die "k6 not found: ${K6_PATH}. Install from https://k6.io/docs/get-started/installation/"
    fi
    verbose "k6 found: $(command -v "$K6_PATH") ($("$K6_PATH" version 2>/dev/null | head -1))"
}

require_jq() {
    if ! command -v jq &>/dev/null; then
        die "jq is required for result parsing"
    fi
}

# ── Format helpers ───────────────────────────────────────────────────
format_duration_ms() {
    local ms="$1"
    if command -v awk &>/dev/null; then
        if awk "BEGIN{exit ($ms >= 1000) ? 0 : 1}" 2>/dev/null; then
            awk "BEGIN{printf \"%.2fs\", $ms / 1000}"
        else
            awk "BEGIN{printf \"%.1fms\", $ms}"
        fi
    else
        echo "${ms}ms"
    fi
}

format_rate() {
    local rate="$1"
    if command -v awk &>/dev/null; then
        awk "BEGIN{printf \"%.1f/s\", $rate}"
    else
        echo "${rate}/s"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          RUN MODE
# ══════════════════════════════════════════════════════════════════════
do_run() {
    [[ -z "$TEST_SCRIPT" ]] && die "No test script specified"
    [[ ! -f "$TEST_SCRIPT" ]] && die "Test script not found: ${TEST_SCRIPT}"
    require_k6

    RUN_ID="$(date +%Y%m%d-%H%M%S)"
    local run_dir="${RESULTS_DIR}/${RUN_ID}"
    mkdir -p "$run_dir"

    section_header "k6 Load Test"
    field "Test script:" "$TEST_SCRIPT"
    field "Virtual users:" "$VUS"
    field "Duration:" "$DURATION"
    field "Run ID:" "$RUN_ID"
    [[ -n "$PUSH_GW" ]] && field "Push gateway:" "$PUSH_GW"
    echo ""

    # Build k6 command
    local -a k6_args=("run")
    k6_args+=("--vus" "$VUS")
    k6_args+=("--duration" "$DURATION")
    k6_args+=("--summary-export" "${run_dir}/summary.json")
    k6_args+=("--out" "json=${run_dir}/results.json")

    if [[ -n "$PUSH_GW" ]]; then
        k6_args+=("--out" "experimental-prometheus-rw=${PUSH_GW}")
    fi

    if [[ -n "$THRESHOLDS_FILE" && -f "$THRESHOLDS_FILE" ]]; then
        k6_args+=("--config" "$THRESHOLDS_FILE")
    fi

    for ev in "${ENV_VARS[@]+"${ENV_VARS[@]}"}"; do
        k6_args+=("-e" "$ev")
    done

    if [[ -n "$TAGS" ]]; then
        k6_args+=("--tag" "$TAGS")
    fi

    k6_args+=("$TEST_SCRIPT")

    verbose "Command: ${K6_PATH} ${k6_args[*]}"

    # Save run metadata
    cat > "${run_dir}/metadata.json" <<EOF
{"run_id":"${RUN_ID}","script":"${TEST_SCRIPT}","vus":${VUS},"duration":"${DURATION}","timestamp":"$(date -u '+%Y-%m-%dT%H:%M:%SZ')"}
EOF

    log "Starting k6 test..."
    echo ""

    local exit_code=0
    "$K6_PATH" "${k6_args[@]}" 2>&1 | tee "${run_dir}/output.log" || exit_code=$?

    echo ""

    # Parse and display results
    if [[ -f "${run_dir}/summary.json" ]]; then
        display_results "${run_dir}/summary.json"
    else
        warn "No summary file generated"
    fi

    section_header "Run Summary"
    field "Run ID:" "$RUN_ID"
    field "Results:" "$run_dir"
    field "Duration:" "$(elapsed)"

    if [[ $exit_code -ne 0 ]]; then
        field_color "Status:" "${RED}FAILED (exit ${exit_code})${RESET}"
    else
        field_color "Status:" "${GREEN}PASSED${RESET}"
    fi

    return "$exit_code"
}

display_results() {
    local summary_file="$1"

    if ! command -v jq &>/dev/null; then
        warn "jq not available — skipping detailed results"
        return
    fi

    section_header "Test Results"

    local http_reqs http_req_dur_avg http_req_dur_p95 http_req_dur_p99
    local http_req_failed iterations data_received data_sent

    http_reqs=$(jq -r '.metrics.http_reqs.values.count // 0' "$summary_file" 2>/dev/null || echo 0)
    http_req_dur_avg=$(jq -r '.metrics.http_req_duration.values.avg // 0' "$summary_file" 2>/dev/null || echo 0)
    http_req_dur_p95=$(jq -r '.metrics.http_req_duration.values["p(95)"] // 0' "$summary_file" 2>/dev/null || echo 0)
    http_req_dur_p99=$(jq -r '.metrics.http_req_duration.values["p(99)"] // 0' "$summary_file" 2>/dev/null || echo 0)
    http_req_failed=$(jq -r '.metrics.http_req_failed.values.rate // 0' "$summary_file" 2>/dev/null || echo 0)
    iterations=$(jq -r '.metrics.iterations.values.count // 0' "$summary_file" 2>/dev/null || echo 0)
    data_received=$(jq -r '.metrics.data_received.values.count // 0' "$summary_file" 2>/dev/null || echo 0)
    data_sent=$(jq -r '.metrics.data_sent.values.count // 0' "$summary_file" 2>/dev/null || echo 0)

    local iter_rate
    iter_rate=$(jq -r '.metrics.iterations.values.rate // 0' "$summary_file" 2>/dev/null || echo 0)

    printf "  ${BOLD}%-28s %s${RESET}\n" "METRIC" "VALUE"
    printf "  %s\n" "$(printf '%.0s─' {1..50})"
    printf "  %-28s %s\n" "HTTP Requests" "$http_reqs"
    printf "  %-28s %s\n" "Request Duration (avg)" "$(format_duration_ms "$http_req_dur_avg")"
    printf "  %-28s %s\n" "Request Duration (p95)" "$(format_duration_ms "$http_req_dur_p95")"
    printf "  %-28s %s\n" "Request Duration (p99)" "$(format_duration_ms "$http_req_dur_p99")"

    local fail_pct
    fail_pct=$(awk "BEGIN{printf \"%.2f\", $http_req_failed * 100}" 2>/dev/null || echo "0")
    if awk "BEGIN{exit ($http_req_failed > 0) ? 0 : 1}" 2>/dev/null; then
        printf "  %-28s ${RED}%s%%${RESET}\n" "Failed Requests" "$fail_pct"
    else
        printf "  %-28s ${GREEN}%s%%${RESET}\n" "Failed Requests" "$fail_pct"
    fi

    printf "  %-28s %s\n" "Iterations" "$iterations"
    printf "  %-28s %s\n" "Iteration Rate" "$(format_rate "$iter_rate")"
    printf "  %-28s %s\n" "Data Received" "$(numfmt --to=iec "$data_received" 2>/dev/null || echo "${data_received} B")"
    printf "  %-28s %s\n" "Data Sent" "$(numfmt --to=iec "$data_sent" 2>/dev/null || echo "${data_sent} B")"
}

# ══════════════════════════════════════════════════════════════════════
#                          LIST MODE
# ══════════════════════════════════════════════════════════════════════
do_list() {
    local dir="${LIST_DIR:-.}"
    [[ ! -d "$dir" ]] && die "Directory not found: ${dir}"

    section_header "Available Test Scripts"

    local count=0
    printf "  ${BOLD}%-40s %10s %s${RESET}\n" "SCRIPT" "SIZE" "MODIFIED"
    printf "  %s\n" "$(printf '%.0s─' {1..65})"

    while IFS= read -r -d '' f; do
        local name size modified
        name=$(basename "$f")
        size=$(stat --printf="%s" "$f" 2>/dev/null || stat -f%z "$f" 2>/dev/null || echo 0)
        modified=$(stat --printf="%y" "$f" 2>/dev/null | cut -d' ' -f1 || stat -f"%Sm" -t "%Y-%m-%d" "$f" 2>/dev/null || echo "unknown")
        local human_size
        human_size=$(numfmt --to=iec "$size" 2>/dev/null || echo "${size}B")
        printf "  %-40s %10s %s\n" "${name:0:38}" "$human_size" "$modified"
        ((count++)) || true
    done < <(find "$dir" -maxdepth 2 -name '*.js' -type f -print0 2>/dev/null | sort -z)

    echo ""
    field "Total scripts:" "$count"

    if [[ "$count" -eq 0 ]]; then
        warn "No .js test scripts found in ${dir}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          COMPARE MODE
# ══════════════════════════════════════════════════════════════════════
do_compare() {
    [[ -z "$COMPARE_RUN" ]] && die "No run ID specified for comparison"
    require_jq

    local prev_summary="${RESULTS_DIR}/${COMPARE_RUN}/summary.json"
    [[ ! -f "$prev_summary" ]] && die "Previous run not found: ${prev_summary}"

    # Find most recent run (excluding the comparison target)
    local latest_dir=""
    while IFS= read -r d; do
        local base
        base=$(basename "$d")
        if [[ "$base" != "$COMPARE_RUN" && -f "${d}/summary.json" ]]; then
            latest_dir="$d"
            break
        fi
    done < <(find "$RESULTS_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort -r)

    if [[ -z "$latest_dir" ]]; then
        die "No other runs found to compare against"
    fi

    local curr_summary="${latest_dir}/summary.json"
    local curr_id
    curr_id=$(basename "$latest_dir")

    section_header "Run Comparison"
    field "Current run:" "$curr_id"
    field "Previous run:" "$COMPARE_RUN"
    echo ""

    printf "  ${BOLD}%-24s %14s %14s %10s${RESET}\n" "METRIC" "CURRENT" "PREVIOUS" "DELTA"
    printf "  %s\n" "$(printf '%.0s─' {1..66})"

    local metrics=("http_reqs.values.count" "http_req_duration.values.avg" "http_req_duration.values[\"p(95)\"]" "iterations.values.count")
    local labels=("HTTP Requests" "Duration (avg ms)" "Duration (p95 ms)" "Iterations")

    for i in "${!metrics[@]}"; do
        local metric="${metrics[$i]}"
        local label="${labels[$i]}"
        local curr_val prev_val
        curr_val=$(jq -r ".metrics.${metric} // 0" "$curr_summary" 2>/dev/null || echo 0)
        prev_val=$(jq -r ".metrics.${metric} // 0" "$prev_summary" 2>/dev/null || echo 0)

        local delta delta_pct color
        delta=$(awk "BEGIN{printf \"%.1f\", $curr_val - $prev_val}" 2>/dev/null || echo "0")
        if awk "BEGIN{exit ($prev_val > 0) ? 0 : 1}" 2>/dev/null; then
            delta_pct=$(awk "BEGIN{printf \"%.1f%%\", (($curr_val - $prev_val) / $prev_val) * 100}" 2>/dev/null || echo "0%")
        else
            delta_pct="N/A"
        fi

        color="$RESET"
        if awk "BEGIN{exit ($delta > 0) ? 0 : 1}" 2>/dev/null; then
            color="$RED"
        elif awk "BEGIN{exit ($delta < 0) ? 0 : 1}" 2>/dev/null; then
            color="$GREEN"
        fi

        printf "  %-24s %14s %14s ${color}%10s${RESET}\n" "$label" \
            "$(awk "BEGIN{printf \"%.1f\", $curr_val}")" \
            "$(awk "BEGIN{printf \"%.1f\", $prev_val}")" \
            "$delta_pct"
    done
}

# ══════════════════════════════════════════════════════════════════════
#                          REPORT MODE
# ══════════════════════════════════════════════════════════════════════
do_report() {
    [[ ! -d "$RESULTS_DIR" ]] && die "Results directory not found: ${RESULTS_DIR}"

    section_header "Test Run History"

    printf "  ${BOLD}%-20s %-30s %8s %12s %s${RESET}\n" "RUN ID" "SCRIPT" "VUS" "DURATION" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..80})"

    local count=0
    while IFS= read -r d; do
        local meta="${d}/metadata.json"
        [[ ! -f "$meta" ]] && continue

        local run_id script vus duration
        run_id=$(basename "$d")
        if command -v jq &>/dev/null; then
            script=$(jq -r '.script // "unknown"' "$meta" 2>/dev/null || echo "unknown")
            vus=$(jq -r '.vus // "?"' "$meta" 2>/dev/null || echo "?")
            duration=$(jq -r '.duration // "?"' "$meta" 2>/dev/null || echo "?")
        else
            script="(jq required)"
            vus="?"
            duration="?"
        fi

        local status_icon="${GREEN}✓${RESET}"
        if [[ -f "${d}/summary.json" ]]; then
            local fail_rate
            fail_rate=$(jq -r '.metrics.http_req_failed.values.rate // 0' "${d}/summary.json" 2>/dev/null || echo 0)
            if awk "BEGIN{exit ($fail_rate > 0) ? 0 : 1}" 2>/dev/null; then
                status_icon="${RED}✗${RESET}"
            fi
        else
            status_icon="${YELLOW}?${RESET}"
        fi

        printf "  %-20s %-30s %8s %12s %b\n" "$run_id" "${script:0:28}" "$vus" "$duration" "$status_icon"
        ((count++)) || true
    done < <(find "$RESULTS_DIR" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort -r)

    echo ""
    field "Total runs:" "$count"

    if [[ "$count" -eq 0 ]]; then
        warn "No test runs found in ${RESULTS_DIR}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          INSPECT MODE
# ══════════════════════════════════════════════════════════════════════
do_inspect() {
    [[ -z "$INSPECT_SCRIPT" ]] && die "No test script specified"
    [[ ! -f "$INSPECT_SCRIPT" ]] && die "Test script not found: ${INSPECT_SCRIPT}"

    section_header "Test Inspection"
    field "Script:" "$INSPECT_SCRIPT"
    field "Size:" "$(stat --printf="%s" "$INSPECT_SCRIPT" 2>/dev/null || stat -f%z "$INSPECT_SCRIPT" 2>/dev/null || echo unknown) bytes"
    echo ""

    log "k6 command that would execute:"
    echo ""
    echo -e "  ${DIM}${K6_PATH} run \\"
    echo -e "    --vus ${VUS} \\"
    echo -e "    --duration ${DURATION} \\"
    [[ -n "$PUSH_GW" ]] && echo -e "    --out experimental-prometheus-rw=${PUSH_GW} \\"
    [[ -n "$THRESHOLDS_FILE" ]] && echo -e "    --config ${THRESHOLDS_FILE} \\"
    for ev in "${ENV_VARS[@]+"${ENV_VARS[@]}"}"; do
        echo -e "    -e ${ev} \\"
    done
    echo -e "    ${INSPECT_SCRIPT}${RESET}"

    echo ""
    if command -v "$K6_PATH" &>/dev/null; then
        field_color "k6 version:" "${GREEN}$("$K6_PATH" version 2>/dev/null | head -1)${RESET}"
    else
        field_color "k6 status:" "${RED}not installed${RESET}"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                          HELP
# ══════════════════════════════════════════════════════════════════════
show_help() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} — k6 Load Test Runner

Run k6 load tests with Prometheus metrics push, result comparison,
and formatted reporting.

${BOLD}MODES${RESET}
  --run SCRIPT         Execute a k6 test script
  --list [DIR]         List available test scripts (default: current dir)
  --compare RUN_ID     Compare latest run against a previous run
  --report             Show test run history
  --inspect SCRIPT     Show what would execute without running

${BOLD}OPTIONS${RESET}
  --vus N              Number of virtual users (default: 10)
  --duration TIME      Test duration (default: 30s)
  --push-gw URL        Prometheus push gateway URL
  --results-dir DIR    Results directory (default: ./k6-results)
  --thresholds FILE    k6 config file with threshold definitions
  --env KEY=VALUE      Pass environment variable to k6 (repeatable)
  --tag KEY=VALUE      Add tag to k6 metrics
  --format FORMAT      Output format: text, json, csv (default: text)
  --verbose            Debug output
  --no-color           Disable colored output
  --help               Show this help message

${BOLD}ENVIRONMENT VARIABLES${RESET}
  K6_PATH              Path to k6 binary (default: k6)
  K6_VUS               Default virtual users (default: 10)
  K6_DURATION           Default test duration (default: 30s)
  K6_PUSH_GATEWAY      Prometheus push gateway URL
  K6_RESULTS_DIR       Results directory (default: ./k6-results)
  K6_FORMAT            Output format (default: text)
  K6_TAGS              Default tags for k6 metrics
  VERBOSE              Enable verbose output (true/false)
  COLOR                Color mode: auto, always, never

${BOLD}EXAMPLES${RESET}
  # Run a load test
  ${SCRIPT_NAME} --run ./tests/load.js

  # Run with custom VUs and duration
  ${SCRIPT_NAME} --run ./tests/api.js --vus 50 --duration 2m

  # Run with Prometheus push
  ${SCRIPT_NAME} --run ./tests/load.js --push-gw http://pushgw:9091

  # Compare runs
  ${SCRIPT_NAME} --compare 20260410-143000

  # List available tests
  ${SCRIPT_NAME} --list ./tests/

  # Inspect without running
  ${SCRIPT_NAME} --inspect ./tests/load.js --vus 100 --duration 5m

${BOLD}EXIT CODES${RESET}
  0   Test passed
  1   Runtime error
  2   Test thresholds exceeded
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                          PARSE ARGS
# ══════════════════════════════════════════════════════════════════════
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --run)        RUN_MODE="run"; TEST_SCRIPT="${2:?--run requires a script path}"; shift 2 ;;
            --list)       RUN_MODE="list"; LIST_DIR="${2:-.}"; shift; if [[ "${1:-}" != -* ]]; then shift; fi ;;
            --compare)    RUN_MODE="compare"; COMPARE_RUN="${2:?--compare requires a run ID}"; shift 2 ;;
            --report)     RUN_MODE="report"; shift ;;
            --inspect)    RUN_MODE="inspect"; INSPECT_SCRIPT="${2:?--inspect requires a script path}"; shift 2 ;;
            --vus)        VUS="${2:?--vus requires a number}"; shift 2 ;;
            --duration)   DURATION="${2:?--duration requires a value}"; shift 2 ;;
            --push-gw)    PUSH_GW="${2:?--push-gw requires a URL}"; shift 2 ;;
            --results-dir) RESULTS_DIR="${2:?--results-dir requires a path}"; shift 2 ;;
            --thresholds) THRESHOLDS_FILE="${2:?--thresholds requires a file}"; shift 2 ;;
            --env)        ENV_VARS+=("${2:?--env requires KEY=VALUE}"); shift 2 ;;
            --tag)        TAGS="${2:?--tag requires KEY=VALUE}"; shift 2 ;;
            --format)     OUTPUT_FORMAT="${2:?--format requires a value}"; export OUTPUT_FORMAT; shift 2 ;;
            --verbose)    VERBOSE="true"; shift ;;
            --no-color)   COLOR="never"; shift ;;
            --help|-h)    setup_colors; show_help; exit 0 ;;
            *)            die "Unknown option: $1 (see --help)" ;;
        esac
    done
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════
main() {
    parse_args "$@"
    setup_colors

    if [[ -z "$RUN_MODE" ]]; then
        err "No mode specified"
        echo ""
        show_help
        exit 1
    fi

    START_TIME=$(date +%s)

    case "$RUN_MODE" in
        run)     do_run ;;
        list)    do_list ;;
        compare) do_compare ;;
        report)  do_report ;;
        inspect) do_inspect ;;
        *)       die "Unknown mode: ${RUN_MODE}" ;;
    esac
}

main "$@"