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

source /usr/local/lib/pia.sh

PIA_IFACE="${PIA_IFACE:-wg0}"
PIA_STATE_DIR="${PIA_STATE_DIR:-/pia}"
PIA_PF_ENABLED="${PIA_PORT_FORWARD:-false}"
PIA_PF_FILE="${PIA_PF_FILE:-$PIA_STATE_DIR/port_forward.json}"
PIA_PF_PORT_FILE="${PIA_PF_PORT_FILE:-$PIA_STATE_DIR/port_forward.port}"
PIA_DNS_ENABLE="${PIA_DNS_ENABLE:-true}"
PIA_DNS_FALLBACK="${PIA_DNS_FALLBACK:-10.0.0.242}"
PIA_HEALTH_URL="${PIA_HEALTH_URL:-https://api64.ipify.org}"
PIA_HANDSHAKE_MAX_AGE_SECONDS="${PIA_HANDSHAKE_MAX_AGE_SECONDS:-180}"
PIA_RECONNECT_BACKOFF_SECONDS="${PIA_RECONNECT_BACKOFF_SECONDS:-5}"

CA_CERT_PATH="/opt/pia/ca.rsa.4096.crt"

pf_pid=""
current_region=""
current_wg_ip=""
current_wg_cn=""

cleanup() {
  set +e
  if [[ -n "$pf_pid" ]]; then
    kill "$pf_pid" 2>/dev/null || true
    wait "$pf_pid" 2>/dev/null || true
    pf_pid=""
  fi
  iptables -D OUTPUT -j PIA_KILLSWITCH 2>/dev/null || true
  iptables -F PIA_KILLSWITCH 2>/dev/null || true
  iptables -X PIA_KILLSWITCH 2>/dev/null || true
  wg-quick down "$PIA_IFACE" 2>/dev/null || true
}

trap cleanup EXIT INT TERM

ensure_state_dir() {
  mkdir -p "$PIA_STATE_DIR"
  chmod 0755 "$PIA_STATE_DIR"
}

write_resolv_conf() {
  local dns="$1"
  log "Setting /etc/resolv.conf to DNS $dns"
  printf "nameserver %s\n" "$dns" > /etc/resolv.conf
}

setup_killswitch() {
  local endpoint_ip="$1"
  local endpoint_port="$2"

  iptables -N PIA_KILLSWITCH 2>/dev/null || true
  iptables -F PIA_KILLSWITCH

  iptables -A PIA_KILLSWITCH -o lo -j ACCEPT
  iptables -A PIA_KILLSWITCH -o "$PIA_IFACE" -j ACCEPT
  iptables -A PIA_KILLSWITCH -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

  # Allow negotiating the tunnel itself (endpoint traffic leaves on eth0)
  iptables -A PIA_KILLSWITCH -p udp -d "$endpoint_ip" --dport "$endpoint_port" -j ACCEPT

  # Allow local traffic (docker internal + loopback ranges)
  iptables -A PIA_KILLSWITCH -d 127.0.0.0/8 -j ACCEPT
  iptables -A PIA_KILLSWITCH -d 172.16.0.0/12 -j ACCEPT
  iptables -A PIA_KILLSWITCH -d 192.168.0.0/16 -j ACCEPT
  iptables -A PIA_KILLSWITCH -d 10.0.0.0/8 -j ACCEPT

  iptables -A PIA_KILLSWITCH -j REJECT

  iptables -C OUTPUT -j PIA_KILLSWITCH 2>/dev/null || iptables -I OUTPUT 1 -j PIA_KILLSWITCH
}

is_tunnel_healthy() {
  if ! wg show "$PIA_IFACE" >/dev/null 2>&1; then
    return 1
  fi

  local now latest age
  now="$(date +%s)"
  latest="$(wg show "$PIA_IFACE" latest-handshakes 2>/dev/null | awk '{print $2}' | sort -nr | head -n1 || true)"
  if [[ -z "$latest" || "$latest" == "0" ]]; then
    return 1
  fi

  age="$(( now - latest ))"
  if (( age > PIA_HANDSHAKE_MAX_AGE_SECONDS )); then
    return 1
  fi

  curl -fsS --max-time 5 "$PIA_HEALTH_URL" >/dev/null
}

start_port_forwarding_loop() {
  local token="$1"
  local wg_ip="$2"
  local wg_cn="$3"
  local region="$4"

  (
    set -euo pipefail

    local sigresp status payload signature port expires_at bindresp

    while true; do
      sigresp="$(pia_pf_get_signature "$wg_ip" "$wg_cn" "$token" "$CA_CERT_PATH")"
      status="$(jq -r '.status' <<<"$sigresp")"
      if [[ "$status" != "OK" ]]; then
        >&2 echo "PF getSignature did not return OK: $sigresp"
        sleep 5
        continue
      fi

      payload="$(jq -r '.payload' <<<"$sigresp")"
      signature="$(jq -r '.signature' <<<"$sigresp")"
      port="$(echo "$payload" | base64 -d | jq -r '.port')"
      expires_at="$(echo "$payload" | base64 -d | jq -r '.expires_at')"

      bindresp="$(pia_pf_bind_port "$wg_ip" "$wg_cn" "$payload" "$signature" "$CA_CERT_PATH")"
      if [[ "$(jq -r '.status' <<<"$bindresp")" != "OK" ]]; then
        >&2 echo "PF bindPort did not return OK: $bindresp"
        sleep 10
        continue
      fi

      ensure_state_dir
      jq -n \
        --arg region "$region" \
        --arg server_ip "$wg_ip" \
        --arg server_cn "$wg_cn" \
        --argjson port "$port" \
        --arg expires_at "$expires_at" \
        --arg refreshed_at "$(date -Is)" \
        '{region:$region,server:{ip:$server_ip,cn:$server_cn},port:$port,expires_at:$expires_at,refreshed_at:$refreshed_at}' \
        > "$PIA_PF_FILE.tmp" \
        && mv "$PIA_PF_FILE.tmp" "$PIA_PF_FILE"

      echo "$port" > "$PIA_PF_PORT_FILE.tmp" && mv "$PIA_PF_PORT_FILE.tmp" "$PIA_PF_PORT_FILE"

      sleep 900
    done
  ) &

  pf_pid="$!"
  log "Port-forward loop started (pid=$pf_pid)."
}

connect_region() {
  local region="$1"
  local wg_ip="$2"
  local wg_cn="$3"
  local token="$4"

  log "Connecting region=$region wg=${wg_cn} (${wg_ip})"

  local privKey pubKey wgjson status peer_ip server_key server_port dns_server
  privKey="$(wg genkey)"
  pubKey="$(echo "$privKey" | wg pubkey)"

  wgjson="$(pia_wireguard_add_key "$wg_ip" "$wg_cn" "$token" "$pubKey" "$CA_CERT_PATH")"
  status="$(jq -r '.status' <<<"$wgjson")"
  if [[ "$status" != "OK" ]]; then
    >&2 echo "WireGuard addKey failed: $wgjson"
    return 1
  fi

  peer_ip="$(jq -r '.peer_ip' <<<"$wgjson")"
  server_key="$(jq -r '.server_key' <<<"$wgjson")"
  server_port="$(jq -r '.server_port' <<<"$wgjson")"
  dns_server="$(jq -r '.dns_servers[0] // empty' <<<"$wgjson")"
  if [[ -z "$dns_server" ]]; then
    dns_server="$PIA_DNS_FALLBACK"
  fi

  mkdir -p /etc/wireguard
  cat >"/etc/wireguard/${PIA_IFACE}.conf" <<EOF
[Interface]
Address = ${peer_ip}
PrivateKey = ${privKey}

[Peer]
PersistentKeepalive = 25
PublicKey = ${server_key}
AllowedIPs = 0.0.0.0/0
Endpoint = ${wg_ip}:${server_port}
EOF

  if [[ "$PIA_DNS_ENABLE" == "true" ]]; then
    write_resolv_conf "$dns_server"
  fi

  wg-quick down "$PIA_IFACE" 2>/dev/null || true
  wg-quick up "$PIA_IFACE"

  setup_killswitch "$wg_ip" "$server_port"

  current_region="$region"
  current_wg_ip="$wg_ip"
  current_wg_cn="$wg_cn"

  log "Connected. Verifying egress."
  curl -fsS --max-time 10 "$PIA_HEALTH_URL" | tr -d '\n' | sed 's/^/PublicIP=/' || true
  echo

  if [[ "$PIA_PF_ENABLED" == "true" ]]; then
    start_port_forwarding_loop "$token" "$wg_ip" "$wg_cn" "$region"
  fi
}

main() {
  require_env PIA_USER
  require_env PIA_PASS
  require_env PIA_REGIONS

  ensure_state_dir
  log "Authenticating to PIA."
  local token
  token="$(pia_get_token)"

  log "Fetching serverlist."
  local serverlist
  serverlist="$(pia_serverlist_v4)"

  local candidates
  candidates="$(pia_pick_region_candidates <<<"$serverlist")"
  if [[ -z "$candidates" ]]; then
    >&2 echo "No matching WireGuard servers found for PIA_REGIONS=$(printf %q "${PIA_REGIONS}")"
    exit 4
  fi

  log "Starting connect loop."
  local line region wg_ip wg_cn connected=0
  while IFS= read -r line; do
    region="${line%%|*}"
    wg_ip="$(cut -d'|' -f2 <<<"$line")"
    wg_cn="$(cut -d'|' -f3 <<<"$line")"

    if connect_region "$region" "$wg_ip" "$wg_cn" "$token"; then
      connected=1
      break
    fi
  done <<<"$candidates"

  if (( connected == 0 )); then
    >&2 echo "Failed to connect to any region in PIA_REGIONS."
    exit 5
  fi

  # Recovery loop (runs forever)
  local idx=0 total
  total="$(wc -l <<<"$candidates" | tr -d ' ')"

  while true; do
    sleep 15
    if is_tunnel_healthy; then
      continue
    fi

    log "Tunnel unhealthy. Reconnecting (rotate region)."
    cleanup

    idx=$(( (idx + 1) % total ))
    line="$(sed -n "$((idx + 1))p" <<<"$candidates")"
    region="${line%%|*}"
    wg_ip="$(cut -d'|' -f2 <<<"$line")"
    wg_cn="$(cut -d'|' -f3 <<<"$line")"

    until connect_region "$region" "$wg_ip" "$wg_cn" "$token"; do
      log "Reconnect failed; sleeping ${PIA_RECONNECT_BACKOFF_SECONDS}s."
      sleep "$PIA_RECONNECT_BACKOFF_SECONDS"
      idx=$(( (idx + 1) % total ))
      line="$(sed -n "$((idx + 1))p" <<<"$candidates")"
      region="${line%%|*}"
      wg_ip="$(cut -d'|' -f2 <<<"$line")"
      wg_cn="$(cut -d'|' -f3 <<<"$line")"
    done
  done
}

main "$@"