#!/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

# Handle --version early, before checking dependencies
if [[ $# -gt 0 ]] && [[ "$1" == "--version" || "$1" == "-v" ]]; then
  CHANGELOG="/usr/share/doc/sensor-control/changelog.gz"
  if [[ -f "${CHANGELOG}" ]]; then
    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

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 'Usage: sr <command> [--sensor <type>] [--sensor_id <id>] [--with-errors]\n' >&2
  printf 'Commands: list, all, internal, external, <sensor-name>\n' >&2
  exit 2
fi

# Parse arguments
SCRIPT=""
FILTER_SENSOR=""
FILTER_SENSOR_ID=""
FILTER_WITH_ERRORS=false
MOCK_MODE=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    --version|-v)
      # Try to read version from debian/changelog
      CHANGELOG="/usr/share/doc/sensor-control/changelog.gz"
      if [[ -f "${CHANGELOG}" ]]; then
        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
      ;;
    --sensor)
      if [[ -z "${2:-}" ]]; then
        printf 'Error: --sensor requires a value\n' >&2
        exit 2
      fi
      FILTER_SENSOR="$2"
      shift 2
      ;;
    --sensor_id)
      if [[ -z "${2:-}" ]]; then
        printf 'Error: --sensor_id requires a value\n' >&2
        exit 2
      fi
      FILTER_SENSOR_ID="$2"
      shift 2
      ;;
    --with-errors)
      FILTER_WITH_ERRORS=true
      shift
      ;;
    --mock)
      MOCK_MODE=true
      shift
      ;;
    -*)
      printf 'Error: Unknown option "%s"\n' "$1" >&2
      exit 2
      ;;
    *)
      if [[ -z "${SCRIPT}" ]]; then
        SCRIPT="$1"
      fi
      shift
      ;;
  esac
done

if [[ -z "${SCRIPT}" ]]; then
  # Default to 'all' if filters are specified but no command given
  if [[ -n "${FILTER_SENSOR}" ]] || [[ -n "${FILTER_SENSOR_ID}" ]] || [[ "${FILTER_WITH_ERRORS}" == "true" ]]; then
    SCRIPT="all"
  else
    printf 'Usage: sr <command> [--sensor <type>] [--sensor_id <id>] [--with-errors] [--mock]\n' >&2
    printf 'Commands: list, all, internal, external, <sensor-name>\n' >&2
    exit 2
  fi
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

readonly LOG_DIR=/var/log/sensor-control
readonly SENS_DIR=/usr/bin
readonly TIMEOUT_SECS=10s
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
      timeout --kill-after="${KILL_AFTER}" "${TIMEOUT_SECS}" "${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

# Apply optional filters to JSON output
apply_filters() {
  local json="$1"
  
  if [[ -n "${FILTER_SENSOR}" ]]; then
    json=$(jq -c --arg sensor "${FILTER_SENSOR}" '[.[] | select(.sensor == $sensor)]' <<<"${json}")
  fi
  
  if [[ -n "${FILTER_SENSOR_ID}" ]]; then
    json=$(jq -c --arg sensor_id "${FILTER_SENSOR_ID}" '[.[] | select(.sensor_id == $sensor_id)]' <<<"${json}")
  fi
  
  if [[ "${FILTER_WITH_ERRORS}" == "true" ]]; then
    json=$(jq -c '[.[] | select(.error != null and .error != "")]' <<<"${json}")
  fi
  
  printf '%s\n' "${json}"
}

# Commands that run all sensors
readonly MULTI_SENSOR_CMDS="all external internal"

# Validate sensor name to prevent path traversal attacks
is_multi_cmd=false
for cmd in ${MULTI_SENSOR_CMDS}; do
  if [[ "${SCRIPT}" == "${cmd}" ]]; then
    is_multi_cmd=true
    break
  fi
done

if [[ "${is_multi_cmd}" == "false" ]]; 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" ]] || [[ "${SCRIPT}" == "internal" ]] || [[ "${SCRIPT}" == "external" ]]; then
  # Determine the command to pass to sensor scripts
  if [[ "${MOCK_MODE}" == "true" ]]; then
    SENSOR_CMD="mock"
  elif [[ "${SCRIPT}" == "all" ]]; then
    SENSOR_CMD="all"
  else
    SENSOR_CMD="${SCRIPT}"
  fi

  # 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}")
    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}" "${SENSOR_CMD}" >"${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 (preserve existing timestamp if set)
    SENS_JSON=$(jq -c --arg node "${PI_SERIAL}" --argjson ts "${TSTAMP}" 'map(.node_id = $node | .timestamp = (.timestamp // $ts))' <<<"${SENS_JSON}")
    ALL_JSONS+=("${SENS_JSON}")
  done

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

  # Merge all arrays and apply filters
  RESULT=$(printf '%s\n' "${ALL_JSONS[@]}" | jq -s 'add' -c)
  apply_filters "${RESULT}"
  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

# Determine sensor command (mock or all)
if [[ "${MOCK_MODE}" == "true" ]]; then
  SINGLE_SENSOR_CMD="mock"
else
  SINGLE_SENSOR_CMD="all"
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}" "${SINGLE_SENSOR_CMD}" 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

# Inject node_id and timestamp (preserve existing timestamp if set)
RESULT=$(jq -c --arg node "${PI_SERIAL}" --argjson ts "${TSTAMP}" 'map(.node_id = $node | .timestamp = (.timestamp // $ts))' <<<"${SENS_JSON}")
apply_filters "${RESULT}"

exit 0
