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

log() { echo "[$(date -Is)] $*" >&2; }

require_env() {
  local name="$1"
  if [[ -z "${!name:-}" ]]; then
    >&2 echo "Missing required env var: $name"
    exit 2
  fi
}

parse_regions() {
  # CSV: "netherlands,france,belgium" or JSON: '["netherlands","france"]'
  local raw="${PIA_REGIONS:-}"
  if [[ -z "$raw" ]]; then
    >&2 echo "Missing required env var: PIA_REGIONS"
    exit 2
  fi

  if [[ "$raw" =~ ^\[.*\]$ ]]; then
    jq -r '.[]' <<<"$raw"
    return
  fi

  tr ',' '\n' <<<"$raw" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' | awk 'NF'
}

read_previous_region() {
  local state_path="${ROTATOR_STATE_PATH:-/config/rotator-state.json}"
  if [[ ! -f "$state_path" ]]; then
    return 0
  fi
  jq -r '.region // empty' "$state_path" 2>/dev/null || true
}

pick_random_region() {
  local regions=() previous candidates=()
  while IFS= read -r region; do
    [[ -n "$region" ]] && regions+=("$region")
  done < <(parse_regions)

  if [[ "${#regions[@]}" -eq 0 ]]; then
    >&2 echo "PIA_REGIONS is empty after parsing"
    exit 2
  fi

  previous="$(read_previous_region)"
  if [[ -n "$previous" ]]; then
    for region in "${regions[@]}"; do
      [[ "$region" == "$previous" ]] && continue
      candidates+=("$region")
    done
  fi

  if [[ "${#candidates[@]}" -eq 0 ]]; then
    if [[ -n "$previous" && "${#regions[@]}" -eq 1 ]]; then
      log "Only one region configured; reusing previous: $previous"
    fi
    candidates=("${regions[@]}")
  else
    log "Excluding previous region: $previous"
  fi

  printf '%s\n' "${candidates[@]}" | shuf -n 1
}

parse_additional_restart_containers() {
  # Sidecars to restart after gluetun. CSV or JSON array.
  local raw="${RESTART_CONTAINERS:-}"
  [[ -z "$raw" ]] && return 0

  if [[ "$raw" =~ ^\[.*\]$ ]]; then
    jq -r '.[]' <<<"$raw"
    return
  fi

  tr ',' '\n' <<<"$raw" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' | awk 'NF'
}

parse_restart_containers() {
  # Gluetun is always restarted first, then any RESTART_CONTAINERS (deduplicated).
  local gluetun="${GLUETUN_CONTAINER:-m3u-filter-vpn}"
  local -A seen=()
  local container

  echo "$gluetun"
  seen["$gluetun"]=1

  while IFS= read -r container; do
    [[ -z "$container" ]] && continue
    [[ -n "${seen[$container]+x}" ]] && continue
    echo "$container"
    seen["$container"]=1
  done < <(parse_additional_restart_containers)
}

restart_containers() {
  local container restarted=()
  while IFS= read -r container; do
    [[ -z "$container" ]] && continue
    log "Restarting container: $container"
    if ! docker restart "$container"; then
      >&2 echo "docker restart failed for container=$container"
      exit 1
    fi
    restarted+=("$container")
  done < <(parse_restart_containers)

  if [[ "${#restarted[@]}" -eq 0 ]]; then
    >&2 echo "No containers configured to restart"
    exit 1
  fi
}

write_state() {
  local region="$1"
  local state_path="${ROTATOR_STATE_PATH:-/config/rotator-state.json}"
  local rotated_at restarted_json
  rotated_at="$(date -Is)"
  restarted_json="$(parse_restart_containers | jq -R -s 'split("\n") | map(select(length > 0))')"
  mkdir -p "$(dirname "$state_path")"
  jq -n \
    --arg region "$region" \
    --arg rotated_at "$rotated_at" \
    --arg wg_config "${WG_CONFIG_PATH:-/config/wireguard/wg0.conf}" \
    --argjson restarted_containers "$restarted_json" \
    '{region: $region, rotated_at: $rotated_at, wg_config: $wg_config, restarted_containers: $restarted_containers}' >"$state_path"
  chmod 644 "$state_path"
}

# EX_TEMPFAIL — entrypoint retries after RATE_LIMIT_WAIT_SECONDS
readonly EXIT_RATE_LIMITED=75

generate_wg_config() {
  local region="$1"
  local out="$2"
  local output rc=0

  # Capture stdout+stderr so we can detect PIA 429 rate limits.
  output="$(pia-wg-config -v -r "$region" -o "$out" "$PIA_USER" "$PIA_PASS" 2>&1)" || rc=$?
  printf '%s\n' "$output" >&2

  if ((rc == 0)); then
    return 0
  fi

  if [[ "$output" == *"too_many_attempts"* || "$output" == *"status 429"* ]]; then
    >&2 echo "pia-wg-config rate-limited for region=$region"
    return "$EXIT_RATE_LIMITED"
  fi

  >&2 echo "pia-wg-config failed for region=$region"
  return 1
}

rotate_once() {
  require_env PIA_USER
  require_env PIA_PASS
  require_env PIA_REGIONS

  local region wg_path tmp_dir
  region="$(pick_random_region)"
  wg_path="${WG_CONFIG_PATH:-/config/wireguard/wg0.conf}"

  log "Selected region: $region"

  tmp_dir="$(mktemp -d)"
  trap 'rm -rf "$tmp_dir"' RETURN

  local tmp_conf="${tmp_dir}/wg0.conf"
  log "Generating WireGuard config with pia-wg-config"
  generate_wg_config "$region" "$tmp_conf"

  if ! grep -q '^\[Interface\]' "$tmp_conf"; then
    >&2 echo "Generated config missing [Interface] section"
    exit 1
  fi

  mkdir -p "$(dirname "$wg_path")"
  mv "$tmp_conf" "$wg_path"
  chmod 600 "$wg_path"
  log "Wrote $wg_path"

  restart_containers
  write_state "$region"

  log "Rotation complete for region=$region"
}

rotate_once