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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPT_OWNER="${CAPTURE_USER:-$(stat -c %U "$SCRIPT_DIR")}"
SCRIPT_OWNER_HOME="$(getent passwd "$SCRIPT_OWNER" | cut -d: -f6)"
cd "$SCRIPT_DIR"

ENV_FILE="$SCRIPT_DIR/.env"
declare -A DOTENV_VALUES

trim_whitespace() {
  local value="$1"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"
  printf '%s' "$value"
}

load_env_file() {
  local line key value quote_char

  if [[ ! -f "$ENV_FILE" ]]; then
    echo "Не найден .env: $ENV_FILE" >&2
    exit 1
  fi

  while IFS= read -r line || [[ -n "$line" ]]; do
    line="${line%$'\r'}"
    line="$(trim_whitespace "$line")"

    [[ -z "$line" || "$line" == \#* ]] && continue

    if [[ "$line" == export\ * ]]; then
      line="${line#export }"
    fi

    [[ "$line" == *=* ]] || continue

    key="$(trim_whitespace "${line%%=*}")"
    value="$(trim_whitespace "${line#*=}")"

    if [[ ${#value} -ge 2 ]]; then
      quote_char="${value:0:1}"
      if [[ ( "$quote_char" == "'" || "$quote_char" == '"' ) && "${value: -1}" == "$quote_char" ]]; then
        value="${value:1:${#value}-2}"
      fi
    fi

    DOTENV_VALUES["$key"]="$value"
  done < "$ENV_FILE"
}

get_env_setting() {
  local key="$1"
  local default="${2-}"

  if [[ -n "${!key+x}" ]]; then
    printf '%s' "${!key}"
  elif [[ -n "${DOTENV_VALUES[$key]+x}" ]]; then
    printf '%s' "${DOTENV_VALUES[$key]}"
  else
    printf '%s' "$default"
  fi
}

get_band_setting() {
  local prefix="$1"
  local band="$2"
  local default="${3-}"

  get_env_setting "${prefix}_${band}" "$(get_env_setting "$prefix" "$default")"
}

load_env_file

############################
# НАСТРОЙКИ
############################

BASE_DIR="$(get_env_setting CAPTURE_BASE_DIR "${SCRIPT_OWNER_HOME}/dataset/noise")"
PYTHON_BIN="$(get_env_setting PYTHON_BIN "$SCRIPT_DIR/.venv-sdr/bin/python")"
SCRIPT_PATH="$(get_env_setting SCRIPT_PATH "$SCRIPT_DIR/scripts_nn/data_saver_headless.py")"

RUN_ONCE="$(get_env_setting RUN_ONCE "0")"
CAPTURE_LOG_FILE="$(get_env_setting CAPTURE_LOG_FILE "$BASE_DIR/capture_hourly.log")"

SYSTEMCTL_BIN=(systemctl)
CURRENT_CAPTURE_PID=""
CURRENT_SERVICE_UNITS=()

lim_all=10
PER_FREQ_LIMIT_BYTES=$((1 * 1024 * 1024 * 1024))
TOTAL_LIMIT_BYTES=$((lim_all * 1024 * 1024 * 1024))
CYCLE_SECONDS="$(get_env_setting CAPTURE_CYCLE_SECONDS "1")"

DEFAULT_SAMP_RATE="$(get_env_setting CAPTURE_SAMP_RATE "20e6")"
DEFAULT_SPLIT_SIZE="$(get_env_setting CAPTURE_SPLIT_SIZE "400000")"
DEFAULT_DELAY="$(get_env_setting CAPTURE_DELAY "0.01")"
DEFAULT_RF_GAIN="$(get_env_setting CAPTURE_RF_GAIN "12")"
DEFAULT_IF_GAIN="$(get_env_setting CAPTURE_IF_GAIN "30")"
DEFAULT_BB_GAIN="$(get_env_setting CAPTURE_BB_GAIN "36")"

ORDER=(433)
SUPPORTED_BANDS=(433 750 915 1200 1500 2400 3300 4500 5200 5800)

declare -A SERIAL
declare -A FREQ_HZ
declare -A SERVICE_UNIT
declare -A BACKEND

for band in "${SUPPORTED_BANDS[@]}"; do
  SERIAL["$band"]="$(get_env_setting "hack_${band}" "$(get_env_setting "HACKID_${band}")")"
  FREQ_HZ["$band"]="$(get_env_setting "c_freq_${band}" "$band")e6"
  SERVICE_UNIT["$band"]="dronedetector-sdr-${band}.service"
  BACKEND["$band"]="$(tr '[:upper:]' '[:lower:]' <<<"$(get_band_setting CAPTURE_BACKEND "$band" "hackrf")")"
done

############################
# ВСПОМОГАТЕЛЬНОЕ
############################

log() {
  printf '[%s] %s\n' "$(date '+%F %T')" "$*"
}

systemctl_run() {
  "${SYSTEMCTL_BIN[@]}" "$@"
}

service_exists() {
  local unit="$1"
  systemctl_run show "$unit" >/dev/null 2>&1
}

service_is_active() {
  local unit="$1"
  systemctl_run is-active --quiet "$unit"
}

remember_service_unit() {
  local unit="$1"
  local existing

  for existing in "${CURRENT_SERVICE_UNITS[@]}"; do
    if [[ "$existing" == "$unit" ]]; then
      return 0
    fi
  done

  CURRENT_SERVICE_UNITS+=("$unit")
}

stop_iio_band_service() {
  local band="$1"
  local unit="${SERVICE_UNIT[$band]:-}"

  CURRENT_SERVICE_UNITS=()

  if [[ -z "$unit" ]]; then
    log "Для band=$band не найден service unit"
    return 0
  fi

  if ! service_exists "$unit"; then
    log "Service unit $unit не установлен, пропускаю"
    return 0
  fi

  if service_is_active "$unit"; then
    remember_service_unit "$unit"
  fi

  log "Останавливаю service $unit перед записью band=$band backend=iio"
  systemctl_run stop "$unit"
}

stop_hackrf_band_services() {
  local band="$1"
  local serial="${SERIAL[$band]:-}"
  local other_band unit other_backend other_serial

  if [[ -z "$serial" ]]; then
    log "Для band=$band пустой serial, пропускаю stop/start сервисов"
    return 0
  fi

  CURRENT_SERVICE_UNITS=()
  for other_band in "${!SERVICE_UNIT[@]}"; do
    unit="${SERVICE_UNIT[$other_band]:-}"
    other_backend="${BACKEND[$other_band]:-hackrf}"
    other_serial="${SERIAL[$other_band]:-}"

    if [[ "$other_backend" != "hackrf" || -z "$unit" || -z "$other_serial" || "$other_serial" != "$serial" ]]; then
      continue
    fi

    if ! service_exists "$unit"; then
      log "Service unit $unit не установлен, пропускаю"
      continue
    fi

    if service_is_active "$unit"; then
      remember_service_unit "$unit"
    fi

    log "Останавливаю service $unit перед записью band=$band"
    systemctl_run stop "$unit"
  done
}

stop_band_service() {
  local band="$1"
  local backend="${BACKEND[$band]:-hackrf}"

  case "$backend" in
    hackrf)
      stop_hackrf_band_services "$band"
      ;;
    iio)
      stop_iio_band_service "$band"
      ;;
    *)
      log "Неизвестный backend=$backend для band=$band"
      return 1
      ;;
  esac
}

start_current_service() {
  local unit

  if [[ "${#CURRENT_SERVICE_UNITS[@]}" -eq 0 ]]; then
    return 0
  fi

  for unit in "${CURRENT_SERVICE_UNITS[@]}"; do
    log "Запускаю service $unit после записи"
    systemctl_run start "$unit"
  done
  CURRENT_SERVICE_UNITS=()
}

cleanup_capture() {
  local unit

  if [[ -n "$CURRENT_CAPTURE_PID" ]] && kill -0 "$CURRENT_CAPTURE_PID" 2>/dev/null; then
    log "Останавливаю текущий PID=$CURRENT_CAPTURE_PID при завершении скрипта"
    kill -TERM "$CURRENT_CAPTURE_PID" 2>/dev/null || true
    wait "$CURRENT_CAPTURE_PID" 2>/dev/null || true
  fi
  CURRENT_CAPTURE_PID=""

  if [[ "${#CURRENT_SERVICE_UNITS[@]}" -gt 0 ]]; then
    for unit in "${CURRENT_SERVICE_UNITS[@]}"; do
      log "Восстанавливаю service $unit при завершении скрипта"
      systemctl_run start "$unit" || true
    done
    CURRENT_SERVICE_UNITS=()
  fi
}

on_signal() {
  local sig="$1"
  log "Получен сигнал $sig, завершаю работу"
  exit 130
}

dir_size_bytes() {
  local path="$1"
  if [[ -e "$path" ]]; then
    du -sb "$path" 2>/dev/null | awk '{print $1}'
  else
    echo 0
  fi
}

total_size_bytes() {
  dir_size_bytes "$BASE_DIR"
}

ensure_requirements() {
  local capture_order band backend

  if [[ ! -x "$PYTHON_BIN" ]]; then
    echo "Не найден python: $PYTHON_BIN" >&2
    exit 1
  fi

  if [[ ! -f "$SCRIPT_PATH" ]]; then
    echo "Не найден script: $SCRIPT_PATH" >&2
    exit 1
  fi

  if ! command -v systemctl >/dev/null 2>&1; then
    echo "Не найден systemctl" >&2
    exit 1
  fi

  if [[ "${EUID}" -ne 0 ]]; then
    if command -v sudo >/dev/null 2>&1; then
      SYSTEMCTL_BIN=(sudo systemctl)
    else
      echo "Для управления service нужен root или sudo" >&2
      exit 1
    fi
  fi

  mkdir -p "$BASE_DIR"
  mkdir -p "$(dirname "$CAPTURE_LOG_FILE")"

  capture_order="$(get_env_setting CAPTURE_ORDER)"
  if [[ -n "$capture_order" ]]; then
    ORDER=()
    for band in ${capture_order//,/ }; do
      ORDER+=("$band")
    done
  fi

  if [[ "${#ORDER[@]}" -eq 0 ]]; then
    echo "ORDER пустой. Укажите частоты в ORDER или задайте CAPTURE_ORDER." >&2
    exit 1
  fi

  for band in "${ORDER[@]}"; do
    backend="${BACKEND[$band]:-}"
    if [[ -z "$backend" ]]; then
      echo "Band $band не поддерживается capture_hourly." >&2
      exit 1
    fi

    case "$backend" in
      hackrf|iio)
        ;;
      *)
        echo "Неизвестный backend=$backend для band=$band" >&2
        exit 1
        ;;
    esac

    if [[ "$backend" == "iio" ]]; then
      if ! command -v iio_attr >/dev/null 2>&1; then
        echo "Не найден iio_attr для band=$band" >&2
        exit 1
      fi
      if ! command -v iio_readdev >/dev/null 2>&1; then
        echo "Не найден iio_readdev для band=$band" >&2
        exit 1
      fi
    fi
  done

  log "Рабочая директория: $SCRIPT_DIR"
  log "BASE_DIR: $BASE_DIR"
  log "ORDER: ${ORDER[*]}"
  log "RUN_ONCE: $RUN_ONCE"
}

run_one_freq() {
  local band="$1"
  local backend="${BACKEND[$band]}"
  local serial="${SERIAL[$band]}"
  local freq="${FREQ_HZ[$band]}"
  local ts out_dir log_file pid
  local samp_rate split_size delay
  local rf_gain if_gain bb_gain
  local iio_uri iio_device iio_phy_device iio_i_channel iio_q_channel
  local iio_lo_channel iio_port_select iio_gain_mode iio_hardwaregain
  local iio_timeout_ms iio_settle iio_bandwidth iio_samples_per_read
  local backend_desc
  local -a cmd

  if [[ "$backend" == "hackrf" && -z "$serial" ]]; then
    log "Для band=$band backend=hackrf, но serial пустой, пропускаю"
    return 0
  fi

  ts="$(date '+%F_%H-%M-%S')"
  out_dir="$BASE_DIR/$band/$ts"
  log_file="$out_dir/run.log"

  samp_rate="$(get_band_setting CAPTURE_SAMP_RATE "$band" "$DEFAULT_SAMP_RATE")"
  split_size="$(get_band_setting CAPTURE_SPLIT_SIZE "$band" "$DEFAULT_SPLIT_SIZE")"
  delay="$(get_band_setting CAPTURE_DELAY "$band" "$DEFAULT_DELAY")"

  cmd=(
    "$PYTHON_BIN"
    "$SCRIPT_PATH"
    "--backend" "$backend"
    "--freq" "$freq"
    "--save-dir" "$out_dir"
    "--file-tag" "${band}_"
    "--samp-rate" "$samp_rate"
    "--split-size" "$split_size"
    "--delay" "$delay"
  )

  if [[ "$backend" == "hackrf" ]]; then
    rf_gain="$(get_band_setting CAPTURE_RF_GAIN "$band" "$DEFAULT_RF_GAIN")"
    if_gain="$(get_band_setting CAPTURE_IF_GAIN "$band" "$DEFAULT_IF_GAIN")"
    bb_gain="$(get_band_setting CAPTURE_BB_GAIN "$band" "$DEFAULT_BB_GAIN")"
    cmd+=(
      "--serial" "$serial"
      "--rf-gain" "$rf_gain"
      "--if-gain" "$if_gain"
      "--bb-gain" "$bb_gain"
    )
    backend_desc="serial=$serial"
  else
    iio_uri="$(get_band_setting CAPTURE_IIO_URI "$band" "ip:192.168.2.1")"
    iio_device="$(get_band_setting CAPTURE_IIO_DEVICE "$band" "cf-ad9361-lpc")"
    iio_phy_device="$(get_band_setting CAPTURE_IIO_PHY_DEVICE "$band" "ad9361-phy")"
    iio_i_channel="$(get_band_setting CAPTURE_IIO_I_CHANNEL "$band" "voltage0")"
    iio_q_channel="$(get_band_setting CAPTURE_IIO_Q_CHANNEL "$band" "voltage1")"
    iio_lo_channel="$(get_band_setting CAPTURE_IIO_LO_CHANNEL "$band" "altvoltage0")"
    iio_port_select="$(get_band_setting CAPTURE_IIO_PORT_SELECT "$band" "A_BALANCED")"
    iio_gain_mode="$(get_band_setting CAPTURE_IIO_GAIN_MODE "$band" "slow_attack")"
    iio_hardwaregain="$(get_band_setting CAPTURE_IIO_HARDWAREGAIN "$band")"
    iio_timeout_ms="$(get_band_setting CAPTURE_IIO_TIMEOUT_MS "$band" "4000")"
    iio_settle="$(get_band_setting CAPTURE_SETTLE "$band" "0.12")"
    iio_bandwidth="$(get_band_setting CAPTURE_IIO_BANDWIDTH "$band")"
    iio_samples_per_read="$(get_band_setting CAPTURE_IIO_SAMPLES_PER_READ "$band")"

    cmd+=(
      "--iio-uri" "$iio_uri"
      "--iio-device" "$iio_device"
      "--iio-phy-device" "$iio_phy_device"
      "--iio-i-channel" "$iio_i_channel"
      "--iio-q-channel" "$iio_q_channel"
      "--iio-lo-channel" "$iio_lo_channel"
      "--iio-port-select" "$iio_port_select"
      "--iio-gain-mode" "$iio_gain_mode"
      "--timeout-ms" "$iio_timeout_ms"
      "--settle" "$iio_settle"
    )

    if [[ -n "$iio_hardwaregain" ]]; then
      cmd+=("--iio-hardwaregain" "$iio_hardwaregain")
    fi
    if [[ -n "$iio_bandwidth" ]]; then
      cmd+=("--bandwidth" "$iio_bandwidth")
    fi
    if [[ -n "$iio_samples_per_read" ]]; then
      cmd+=("--samples-per-read" "$iio_samples_per_read")
    fi

    backend_desc="uri=$iio_uri"
  fi

  log "Старт band=$band backend=$backend $backend_desc freq=$freq dir=$out_dir"
  if ! mkdir -p "$out_dir"; then
    log "Не удалось создать каталог $out_dir"
    return 1
  fi

  stop_band_service "$band"

  "${cmd[@]}" >"$log_file" 2>&1 &
  pid=$!
  CURRENT_CAPTURE_PID="$pid"
  log "PID=$pid"

  while kill -0 "$pid" 2>/dev/null; do
    local cur_dir_size cur_total_size
    cur_dir_size="$(dir_size_bytes "$out_dir")"
    cur_total_size="$(total_size_bytes)"

    if (( cur_total_size >= TOTAL_LIMIT_BYTES )); then
      log "Достигнут общий лимит ${lim_all} GiB. Останавливаю PID=$pid"
      kill -TERM "$pid" 2>/dev/null || true
      wait "$pid" 2>/dev/null || true
      CURRENT_CAPTURE_PID=""
      start_current_service
      return 2
    fi

    if (( cur_dir_size >= PER_FREQ_LIMIT_BYTES )); then
      log "Для band=$band достигнут лимит 1 GiB. Останавливаю PID=$pid"
      kill -TERM "$pid" 2>/dev/null || true

      for _ in {1..10}; do
        if ! kill -0 "$pid" 2>/dev/null; then
          break
        fi
        sleep 1
      done

      if kill -0 "$pid" 2>/dev/null; then
        log "PID=$pid не завершился по TERM, отправляю KILL"
        kill -KILL "$pid" 2>/dev/null || true
      fi

      wait "$pid" 2>/dev/null || true
      break
    fi

    sleep 1
  done

  CURRENT_CAPTURE_PID=""
  start_current_service
  log "Завершен band=$band backend=$backend, размер=$(du -sh "$out_dir" | awk '{print $1}')"
  return 0
}

main_loop() {
  local total_before cycle_start elapsed sleep_left rc

  while true; do
    total_before="$(total_size_bytes)"

    if (( total_before >= TOTAL_LIMIT_BYTES )); then
      log "Общий размер уже >= ${lim_all} GiB, выхожу"
      break
    fi

    cycle_start="$(date +%s)"
    log "Новый цикл"

    for band in "${ORDER[@]}"; do
      if (( $(total_size_bytes) >= TOTAL_LIMIT_BYTES )); then
        log "Общий лимит достигнут внутри цикла, выхожу"
        return 0
      fi

      run_one_freq "$band" || {
        rc=$?
        if [[ $rc -eq 2 ]]; then
          log "Остановка по общему лимиту"
          return 0
        fi
        log "Ошибка записи band=$band rc=$rc"
        return "$rc"
      }
    done

    if [[ "$RUN_ONCE" == "1" ]]; then
      log "RUN_ONCE=1, завершаю работу после одного цикла"
      break
    fi

    elapsed=$(( $(date +%s) - cycle_start ))
    sleep_left=$(( CYCLE_SECONDS - elapsed ))

    if (( sleep_left > 0 )); then
      log "Цикл занял ${elapsed} сек, жду ${sleep_left} сек до следующего запуска"
      sleep "$sleep_left"
    else
      log "Цикл занял ${elapsed} сек, паузы нет"
    fi
  done
}

trap cleanup_capture EXIT
trap 'on_signal INT' INT
trap 'on_signal TERM' TERM

ensure_requirements
main_loop