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

YGG_CONF="${YGG_CONF:-/etc/yggdrasil/yggdrasil.conf}"
YGG_CONNECT_WAIT_SECONDS="${YGG_CONNECT_WAIT_SECONDS:-8}"
YGG_RUNTIME_CONF="${YGG_RUNTIME_CONF:-/run/yggdrasil.conf}"
PROXY_CFG_PRIMARY="${PROXY_CFG_PRIMARY:-/etc/3proxy/first-instanse.cfg}"
PROXY_CFG_SECONDARY="${PROXY_CFG_SECONDARY:-/etc/3proxy/second-instanse.cfg}"
SOCKS_ROUTER_BIN="${SOCKS_ROUTER_BIN:-/usr/local/bin/socks_router.py}"

ygg_pid=""
proxy_primary_pid=""
proxy_secondary_pid=""
router_pid=""

log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"
}

cleanup() {
  if [[ -n "${proxy_primary_pid}" ]]; then
    kill -TERM "${proxy_primary_pid}" 2>/dev/null || true
  fi
  if [[ -n "${proxy_secondary_pid}" ]]; then
    kill -TERM "${proxy_secondary_pid}" 2>/dev/null || true
  fi
  if [[ -n "${router_pid}" ]]; then
    kill -TERM "${router_pid}" 2>/dev/null || true
  fi
  if [[ -n "${ygg_pid}" ]]; then
    kill -TERM "${ygg_pid}" 2>/dev/null || true
  fi
  wait || true
}

trap cleanup TERM INT

if [[ ! -f "${YGG_CONF}" ]]; then
  log "Missing Yggdrasil config at ${YGG_CONF}"
  exit 1
fi

if [[ ! -f "${PROXY_CFG_PRIMARY}" ]]; then
  log "Missing primary 3proxy config at ${PROXY_CFG_PRIMARY}"
  exit 1
fi

if [[ ! -f "${PROXY_CFG_SECONDARY}" ]]; then
  log "Missing secondary 3proxy config at ${PROXY_CFG_SECONDARY}"
  exit 1
fi

if [[ ! -f "${SOCKS_ROUTER_BIN}" ]]; then
  log "Missing SOCKS router at ${SOCKS_ROUTER_BIN}"
  exit 1
fi

cp "${YGG_CONF}" "${YGG_RUNTIME_CONF}"

# Populate keys once if config still has empty key fields.
if ! grep -Eq '^[[:space:]]*PrivateKey:[[:space:]]+[0-9a-fA-F]+' "${YGG_RUNTIME_CONF}"; then
  log "Generating Yggdrasil keys for dark3proxy"
  key_output="$(timeout 8 /usr/bin/genkeys 2>/dev/null || true)"
  priv_key="$(printf '%s\n' "${key_output}" | awk '/^Priv:/{print $2; exit}')"
  pub_key="$(printf '%s\n' "${key_output}" | awk '/^Pub:/{print $2; exit}')"

  if [[ -z "${priv_key}" || -z "${pub_key}" ]]; then
    log "Failed to parse generated Yggdrasil keys"
    exit 1
  fi

  sed -i -E \
    -e "s|^([[:space:]]*PublicKey:).*|\1 ${pub_key}|" \
    -e "s|^([[:space:]]*PrivateKey:).*|\1 ${priv_key}|" \
    "${YGG_RUNTIME_CONF}"
fi

sysctl net.ipv6.conf.all.disable_ipv6=0 >/dev/null 2>&1 || true

log "Starting Yggdrasil"
/usr/bin/yggdrasil -useconffile "${YGG_RUNTIME_CONF}" &
ygg_pid=$!

sleep "${YGG_CONNECT_WAIT_SECONDS}"
if ! kill -0 "${ygg_pid}" 2>/dev/null; then
  log "Yggdrasil exited during startup"
  exit 1
fi

log "Starting secondary 3proxy instance (${PROXY_CFG_SECONDARY})"
/usr/bin/3proxy "${PROXY_CFG_SECONDARY}" &
proxy_secondary_pid=$!

sleep 1
if ! kill -0 "${proxy_secondary_pid}" 2>/dev/null; then
  log "Secondary 3proxy instance exited during startup"
  exit 1
fi

log "Starting hostname-preserving SOCKS router (${SOCKS_ROUTER_BIN})"
python3 "${SOCKS_ROUTER_BIN}" &
router_pid=$!

sleep 1
if ! kill -0 "${router_pid}" 2>/dev/null; then
  log "SOCKS router exited during startup"
  exit 1
fi

log "Starting primary 3proxy instance (${PROXY_CFG_PRIMARY})"
/usr/bin/3proxy "${PROXY_CFG_PRIMARY}" &
proxy_primary_pid=$!

sleep 1
if ! kill -0 "${proxy_primary_pid}" 2>/dev/null; then
  log "Primary 3proxy instance exited during startup"
  exit 1
fi

wait -n "${proxy_secondary_pid}" "${proxy_primary_pid}" "${router_pid}" "${ygg_pid}"
exit_code=$?
log "A managed process exited, shutting down"
cleanup
exit "${exit_code}"