#!/bin/bash

set -uo pipefail
trap 'echo "[✖] Execution halted at line $LINENO. Please consult your nearest bash therapist." >&2' ERR
trap 'echo "[✌️] Script exited cleanly. Have a burrito." >&2' EXIT

NO_ENRICH=0
for arg in "$@"; do
  case $arg in
    --no-enrich)
      NO_ENRICH=1
      shift
      ;;
    *)
      echo "Unknown argument: $arg"
      exit 1
      ;;
  esac
done

IFS=$'\n\t'

# === Start logging with 'script' safely ===
if [ -z "${SCRIPT_LOG_STARTED:-}" ]; then
  export SCRIPT_LOG_STARTED=1
  BOOT_ID=$(cat /proc/sys/kernel/random/boot_id)
  LOG_FILE="${HOME}/${BOOT_ID}-session.log"
  echo "[+] Starting session logging to $LOG_FILE"
  exec script -q "$LOG_FILE" --return --flush --command "$0"
fi

# === Now start your real script ===
source settings.env

# (redefine BOOT_ID, TEST_FILE, etc after reloading settings.env if needed)
BOOT_ID=$(cat /proc/sys/kernel/random/boot_id)
TEST_FILE="${HOME}/${BOOT_ID}-speedtest.csv"
ENRICHED_FILE="${TEST_FILE%.csv}+rf.csv"
SSID_METRICS_FILE="${ENRICHED_FILE%.csv}-ssid-metrics.csv"
FAILURE_LOG="${TEST_FILE%.csv}-failures.log"
SCRIPT_START=$(date +%s)

EMAIL_BODY=${EMAIL_BODY:-"Test $BOOT_ID completed successfully. Please find the attached files and collect the probe."}

log() {
  echo "[+] $*" >&2
}

warn() {
  echo "[!] $*" >&2
}

die() {
  warn "$*"
  exit 1
}

[ -z "${RECIPIENT:-}" ] && die "Please set the RECIPIENT variable in settings.env."

log "Test file: $TEST_FILE"
log "Enriched file: $ENRICHED_FILE"
log "SSID metrics file: $SSID_METRICS_FILE"

sudo -v
while true; do sudo -n true; sleep 60; done 2>/dev/null &
SUDO_KEEPALIVE_PID=$!

log "Starting kismet..."
SECONDS=0
sudo systemctl start kismet
log "Kismet startup took $SECONDS seconds"

log "Saturating the capture..."
sleep "$LEAD_TIME"

get_tx_failed() {
  iw dev "$INTERFACE" station dump | awk '/tx failed/ {print $3}'
}

freq_to_channel() {
  local freq=$1
  if (( freq >= 2412 && freq <= 2472 )); then echo $(((freq - 2407) / 5))
  elif (( freq == 2484 )); then echo 14
  elif (( freq >= 5180 && freq <= 5825 )); then echo $(((freq - 5000) / 5))
  else echo "Unknown"; fi
}

run_iperf() {
  local target=$1 mode=$2 direction=$3
  local args=("-c" "$target" "-J" "-t" "$IPERF_DURATION")
  [ "$mode" = "udp" ] && args+=("-u" "-b" "$BANDWIDTH")
  [ "$direction" = "down" ] && args+=("--reverse")

  local tmp_err=$(mktemp) tmp_json=$(mktemp)
  local timeout_duration=$((IPERF_DURATION + 20))

  log "Running iperf3 $mode $direction to $target..."
  SECONDS=0
  if ! timeout "${timeout_duration}s" iperf3 "${args[@]}" >"$tmp_json" 2>"$tmp_err"; then
    warn "iperf3 $mode $direction to $target failed or timed out"
    echo "0"
    return
  fi
  log "iperf3 completed in $SECONDS seconds"

  SECONDS=0
  local parsed
  parsed=$(jq -r '
    if .error then "iperf3-error"
    elif has("end") | not then "no-end"
    elif .end | has("sum_received") then .end.sum_received.bits_per_second
    elif .end | has("sum") then .end.sum.bits_per_second
    else "unexpected-format" end' "$tmp_json" 2>/dev/null || echo "execution-failed")

  log "Parsed iperf3 result in $SECONDS seconds"

  if [[ "$parsed" =~ ^(iperf3-error|no-end|unexpected-format|execution-failed)$ ]]; then
    timestamp=$(date -Iseconds)
    echo "$timestamp,iperf $mode $direction to $target failed with '$parsed'" >>"$FAILURE_LOG"
    echo "[stderr] $(cat "$tmp_err")" >>"$FAILURE_LOG"
    echo "[json] $(cat "$tmp_json")" >>"$FAILURE_LOG"
    echo "0"
  else
    echo "$parsed"
  fi

  rm -f "$tmp_err" "$tmp_json"
}

# Send start email
SECONDS=0
echo -e "Subject: Test ${BOOT_ID} Started\n\nTest ${BOOT_ID} has commenced." | msmtp "$RECIPIENT"
log "Start email sent in $SECONDS seconds"

FAILED_START=$(get_tx_failed)

# CSV setup
[ ! -f "$TEST_FILE" ] && echo "StartTimestamp,EndTimestamp,Link,Level,Noise,BSSID,TX Bitrate,RX Bitrate,$(speedtest --csv-header),TX Failures,Channel,Frequency,Packet Loss,Jitter,LocalTCPUp,LocalTCPDown,LocalUDPUp,LocalUDPDown,RemoteTCPUp,RemoteTCPDown,RemoteUDPUp,RemoteUDPDown" >"$TEST_FILE"

for ((COUNTER = 1; COUNTER <= NUM_TESTS; COUNTER++)); do
  log "Test run $COUNTER of $NUM_TESTS"
  for ((i = 1; i <= NUM_SAMPLES; i++)); do
    log "  Sample $i of $NUM_SAMPLES"
    START_TIME=$(date -Iseconds)

    link_level_noise=$(awk 'NR==3 {gsub(/\./, "", $3); gsub(/\./, "", $4); gsub(/\./, "", $5); print $3","$4","$5}' /proc/net/wireless)
    bssid_and_bitrate=$(iw dev "$INTERFACE" link | awk '/Connected/ {bssid=$3} /tx bitrate/ {tx=$3} /rx bitrate/ {rx=$3} END {print bssid","tx","rx}')

    speed_results=""
    for ((retry = 1; retry <= MAX_RETRIES; retry++)); do
      SECONDS=0
      log "    Speed test attempt $retry"
      speed_results=$(speedtest --secure --csv 2>/dev/null || true)
      log "    Speed test took $SECONDS seconds"
      [[ -n "$speed_results" ]] && break
      warn "    Speedtest failed. Retrying in $RETRY_DELAY seconds..."
      echo "$(date -Iseconds),Speedtest failed on attempt $retry for test $COUNTER, sample $i" >> "$FAILURE_LOG"
      sleep "$RETRY_DELAY"
    done

    if [[ -z "$speed_results" ]]; then
      timestamp=$(date -Iseconds)
      warn "    Speedtest permanently failed. Skipping sample."
      echo "$timestamp,Test $COUNTER,Sample $i" >>"$FAILURE_LOG"
      continue
    fi

    SECONDS=0
    log "    Gathering TX failed count..."
    FAILED_NOW=$(get_tx_failed)
    FAILED_DELTA=$((FAILED_NOW - FAILED_START))
    FAILED_START=$FAILED_NOW
    log "    TX count gathered in $SECONDS seconds"

    freq=$(iw dev "$INTERFACE" link | awk '/freq:/ {print $2}')
    channel=$(freq_to_channel "$freq")

    SECONDS=0
    log "    Running ping test..."
    ping_output=$(ping -c "$PING_COUNT" "$PING_TARGET")
    packet_loss=$(echo "$ping_output" | grep -oP '\d+(?=% packet loss)')
    jitter=$(echo "$ping_output" | grep "time=" | awk '{print $(NF-1)}' | sed 's/time=//g' | awk '{sum+=$1; sumsq+=$1*$1} END {if (NR>1) print sqrt(sumsq/NR - (sum/NR)**2); else print 0}')
    log "    Ping test took $SECONDS seconds"

    log "    Running iperf3 tests..."
    LocalTCPUp=$(run_iperf "$IPERF_LOCAL_TARGET" tcp up)
    LocalTCPDown=$(run_iperf "$IPERF_LOCAL_TARGET" tcp down)
    LocalUDPUp=$(run_iperf "$IPERF_LOCAL_TARGET" udp up)
    LocalUDPDown=$(run_iperf "$IPERF_LOCAL_TARGET" udp down)
    RemoteTCPUp=$(run_iperf "$IPERF_REMOTE_TARGET" tcp up)
    RemoteTCPDown=$(run_iperf "$IPERF_REMOTE_TARGET" tcp down)
    RemoteUDPUp=$(run_iperf "$IPERF_REMOTE_TARGET" udp up)
    RemoteUDPDown=$(run_iperf "$IPERF_REMOTE_TARGET" udp down)

    END_TIME=$(date -Iseconds)
    echo "$START_TIME,$END_TIME,$link_level_noise,$bssid_and_bitrate,$speed_results,$FAILED_DELTA,$channel,$freq,$packet_loss,$jitter,$LocalTCPUp,$LocalTCPDown,$LocalUDPUp,$LocalUDPDown,$RemoteTCPUp,$RemoteTCPDown,$RemoteUDPUp,$RemoteUDPDown" >>"$TEST_FILE"
  done

  [[ "$COUNTER" -lt "$NUM_TESTS" ]] && log "Waiting $TIME_BETWEEN before next test..." && sleep "$TIME_BETWEEN"
done

log "Stopping kismet..."
sudo systemctl stop kismet

if [ "$NO_ENRICH" -eq 0 ]; then
  log "Enriching data..."
  KISMET_LOG=$(find "$KISMET_LOG_DIR" -type f -name "*.pcapng" -printf "%T@ %p\n" | sort -n | tail -1 | cut -d' ' -f2-)
  [ ! -f "$KISMET_LOG" ] && die "Packet capture not found."

  SECONDS=0
  python3 "$SCRIPT_DIRECTORY/enrich.py" --csv "$TEST_FILE" --pcapng "$KISMET_LOG" --output "$ENRICHED_FILE"
  log "Enrichment took $SECONDS seconds"
else
  log "Skipping enrichment as per --no-enrich flag."
fi

ATTACHMENTS=()
[ -f "$ENRICHED_FILE" ] && ATTACHMENTS+=("$ENRICHED_FILE")
[ -f "$FAILURE_LOG" ] && ATTACHMENTS+=("$FAILURE_LOG")
[ -f "$SSID_METRICS_FILE" ] && ATTACHMENTS+=("$SSID_METRICS_FILE")

if [ ${#ATTACHMENTS[@]} -eq 0 ]; then
  warn "No files to attach. Email not sent."
else
  for file in "${ATTACHMENTS[@]}"; do log "Attaching: $file"; done
  echo "$EMAIL_BODY" | mutt -s "Test ${BOOT_ID} Complete" -a "${ATTACHMENTS[@]}" -- "$RECIPIENT"
  log "Email sent to $RECIPIENT with attachments."
fi

sudo kill "$SUDO_KEEPALIVE_PID"
SCRIPT_END=$(date +%s)
log "Full test cycle completed in $((SCRIPT_END - SCRIPT_START)) seconds"