#!/usr/bin/env bash

# WildlifeSystems - sensor read
#
# Standard wrapper for reading sensors in WildlifeSystems. Handles running
# sensor scripts (serially or in parallel), inserting device serial number
# and timestamp into JSON responses, and aggregating results.
#
# For more information: https://docs.wildlife.systems

set -u
set -o pipefail

readonly REQUIRED_CMDS=(jq timeout find basename sed awk pi-data mktemp)

missing=()
for _cmd in "${REQUIRED_CMDS[@]}"; do
  if ! command -v "$_cmd" >/dev/null 2>&1; then
    missing+=("$_cmd")
  fi
done
if (( ${#missing[@]} )); then
  printf 'Missing required commands: %s\n' "${missing[*]}" >&2
  exit 1
fi

if [[ $# -eq 0 ]]; then
  printf 'No arguments supplied\n' >&2
  exit 2
fi

# Handle --version flag
if [[ "$1" == "--version" ]] || [[ "$1" == "-v" ]]; then
  # Try to read version from debian/changelog
  CHANGELOG="/usr/share/doc/sensor-control/changelog.gz"
  if [[ -f "${CHANGELOG}" ]]; then
    # Extract version from first line of changelog
    VERSION=$(zcat "${CHANGELOG}" 2>/dev/null | head -n1 | sed -n 's/^sensor-control (\([^)]*\)).*/\1/p')
    if [[ -n "${VERSION}" ]]; then
      printf 'sr (sensor-control) %s\n' "${VERSION}"
    else
      printf 'sr (sensor-control) version unknown\n'
    fi
  else
    printf 'sr (sensor-control) version unknown (changelog not found)\n'
  fi
  exit 0
fi

SELF="${BASH_SOURCE[0]:-$0}"
if [[ "${SELF}" != /* ]]; then
  if command -v "$(basename "${SELF}")" >/dev/null 2>&1; then
    SELF=$(command -v "$(basename "${SELF}")")
  else
    SELF="$(pwd)/${SELF}"
  fi
fi

SCRIPT="$1"
DEV=${2:-all}

readonly LOG_DIR=/var/log/sensor-control
readonly SENS_DIR=/usr/bin
readonly TIMEOUT_SECS=5
readonly KILL_AFTER=5s

mkdir -p "${LOG_DIR}" 2>/dev/null || true

# Get Pi serial number with validation
if ! PI_SERIAL=$(pi-data serial 2>/dev/null); then
  printf 'Error: Failed to retrieve Pi serial number\n' >&2
  exit 1
fi

if [[ -z "${PI_SERIAL}" ]]; then
  printf 'Error: Pi serial number is empty\n' >&2
  exit 1
fi

TSTAMP=$(date +%s)

run_list() {
  find "${SENS_DIR}" -maxdepth 1 -name 'sensor-*' -executable -print0 |
    while IFS= read -r -d '' SENS_SCRIPT; do
      "${SENS_SCRIPT}" identify 2>/dev/null
      rc=$?
      if [[ $rc -eq 60 ]]; then
        SCRIPT_NAME=$(basename "${SENS_SCRIPT}")
        printf '%s\n' "${SCRIPT_NAME#sensor-}"
      fi
    done
}

if [[ "${SCRIPT}" == "list" ]]; then
  run_list
  exit 0
fi

# Validate sensor name to prevent path traversal attacks
if [[ "${SCRIPT}" != "all" ]] && [[ "${SCRIPT}" != "external" ]] && [[ "${SCRIPT}" != "internal" ]]; then
  if [[ ! "${SCRIPT}" =~ ^[a-zA-Z0-9_-]+$ ]]; then
    printf 'Error: Invalid sensor name "%s"\n' "${SCRIPT}" >&2
    exit 2
  fi
fi

# Default concurrency: from environment or CPU count, fall back to 4
CONCURRENCY=${CONCURRENCY:-$(nproc 2>/dev/null || getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)}
if ! [[ "${CONCURRENCY}" =~ ^[0-9]+$ ]] || [[ ${CONCURRENCY} -lt 1 ]]; then
  CONCURRENCY=4
fi

if [[ "${SCRIPT}" == "all" ]]; then
  # Run sensor scripts in parallel, capture outputs to temp files
  declare -a SENSOR_NAMES=()
  declare -a OUT_FILES=()
  declare -a PIDS=()

  # Cleanup temp files on exit or interruption
  trap 'rm -f "${OUT_FILES[@]}" 2>/dev/null' EXIT INT TERM

  # Launch jobs
  while IFS= read -r SENSOR_NAME; do
    SENSOR_NAMES+=("${SENSOR_NAME}")
    DEV=all
    SENS_SCRIPT="${SENS_DIR}/sensor-${SENSOR_NAME}"
    LOGFILE="${LOG_DIR}/${SENSOR_NAME}.log"

    OUT=$(mktemp)
    OUT_FILES+=("${OUT}")

    # Run sensor script in background, capture stdout to OUT, append stderr to LOGFILE
    timeout --preserve-status --signal=TERM --kill-after=${KILL_AFTER} ${TIMEOUT_SECS} "${SENS_SCRIPT}" "${DEV}" >"${OUT}" 2>>"${LOGFILE}" &
    PIDS+=("$!")

    # Throttle concurrency: if we have reached the limit, wait for the oldest job
    while (( ${#PIDS[@]} >= CONCURRENCY )); do
      pid_to_wait=${PIDS[0]}
      wait "${pid_to_wait}" || true
      # Remove first element
      PIDS=("${PIDS[@]:1}")
    done
  done < <(run_list)

  # Wait for remaining jobs
  for pid in "${PIDS[@]}"; do
    wait "${pid}" || true
  done

  # Process outputs in the same order as SENSOR_NAMES
  declare -a ALL_JSONS=()
  for i in "${!SENSOR_NAMES[@]}"; do
    OUT=${OUT_FILES[i]}

    # Read output safely
    SENS_JSON=""
    if [[ -s "${OUT}" ]]; then
      SENS_JSON=$(<"${OUT}")
    fi

    # Cleanup temp file
    rm -f "${OUT}"

    if [[ -z "${SENS_JSON}" ]]; then
      continue
    fi

    # Validate JSON before processing
    if ! jq -e . >/dev/null 2>&1 <<<"${SENS_JSON}"; then
      printf 'Warning: Invalid JSON from sensor "%s", skipping\n' "${SENSOR_NAMES[i]}" >&2
      continue
    fi

    # Inject node_id and timestamp
    SENS_JSON=$(jq -c --arg node "${PI_SERIAL}" --argjson ts "${TSTAMP}" 'map(.node_id = $node | .timestamp = $ts)' <<<"${SENS_JSON}")
    ALL_JSONS+=("${SENS_JSON}")
  done

  if (( ${#ALL_JSONS[@]} == 0 )); then
    printf '[]\n'
    exit 0
  fi

  # Merge all arrays in one jq invocation
  printf '%s\n' "${ALL_JSONS[@]}" | jq -s 'add' -c
  exit 0
fi

if [[ "${SCRIPT}" == "external" ]]; then
  "${SELF}" all | jq -c 'del(.[] | select(.internal == true))'
  exit 0
fi

if [[ "${SCRIPT}" == "internal" ]]; then
  "${SELF}" all | jq -c 'del(.[] | select(.internal == false))'
  exit 0
fi

# Single sensor invocation (non-parallel)
SENS_SCRIPT="${SENS_DIR}/sensor-${SCRIPT}"

# Check if sensor script exists
if [[ ! -x "${SENS_SCRIPT}" ]]; then
  printf 'Error: Sensor "%s" not found or not executable\n' "${SCRIPT}" >&2
  exit 21
fi

if [[ "${DEV}" == "list" ]]; then
  "${SENS_SCRIPT}" list
  exit 0
fi

SENSOR_BASENAME=$(basename "${SENS_SCRIPT}")
SENS_LOG="${LOG_DIR}/${SENSOR_BASENAME}.log"
rc=0
SENS_JSON=$(timeout --preserve-status --signal=TERM --kill-after=${KILL_AFTER} ${TIMEOUT_SECS} "${SENS_SCRIPT}" "${DEV}" 2>>"${SENS_LOG}") || rc=$?
rc=${rc:-0}

if [[ $rc -ne 0 ]]; then
  if [[ $rc -eq 20 ]]; then
    printf 'Unknown device\n'
  elif [[ $rc -eq 21 ]]; then
    printf 'Unknown sensor\n'
  fi
  exit "$rc"
fi

# Validate JSON output
if [[ -z "${SENS_JSON}" ]]; then
  printf 'Error: Sensor produced no output\n' >&2
  exit 1
fi

if ! jq -e . >/dev/null 2>&1 <<<"${SENS_JSON}"; then
  printf 'Error: Sensor produced invalid JSON\n' >&2
  exit 1
fi

jq -c --arg node "${PI_SERIAL}" --argjson ts "${TSTAMP}" 'map(.node_id = $node | .timestamp = $ts)' <<<"${SENS_JSON}"

exit 0
