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

: "${OC_URL:?OC_URL required}"
: "${OC_SERVERCERT:?OC_SERVERCERT required}"

NOVNC_PORT="${NOVNC_PORT:-6901}"
VNC_PASSWORD="${VNC_PASSWORD:-changeme}"
DISPLAY_ADDR="${DISPLAY:-:1}"
OC_INTERFACE="${OC_INTERFACE:-tun0}"
OC_USER="${OC_USER:-}"
OC_TOTP_SECRET="${OC_TOTP_SECRET:-}"

# Default to hidden browser if OC_USER is set
if [[ -n "$OC_USER" ]]; then
  OC_SSO_ARGS_DEFAULT="--browser-display-mode hidden -u $OC_USER"
else
  OC_SSO_ARGS_DEFAULT="--browser-display-mode shown"
fi

CLOUDFLARED_MODE="${CLOUDFLARED_MODE:-off}"   # off|token|config
CLOUDFLARED_TOKEN="${CLOUDFLARED_TOKEN:-}"
SSH_TUNNEL_ENABLE="${SSH_TUNNEL_ENABLE:-0}"
SSH_DEST="${SSH_DEST:-zawa@10.3.1.201}"
SSH_FORWARDS="${SSH_FORWARDS:-0.0.0.0:8090:localhost:8090}"

pids=()

# Setup keyring with TOTP secret if provided
setup_keyring() {
  if [[ -n "$OC_TOTP_SECRET" && -n "$OC_USER" ]]; then
    python3 -c "
import keyring
keyring.set_password('openconnect-sso', 'totp/$OC_USER', '$OC_TOTP_SECRET'.upper())
print('TOTP secret stored in keyring for $OC_USER')
"
  fi
}

# Create vpn_connect command in PATH and save environment
create_vpn_command() {
  # Save environment variables to a file
  cat > /etc/vpn.env << ENVFILE
export OC_URL="$OC_URL"
export OC_SERVERCERT="$OC_SERVERCERT"
export OC_INTERFACE="$OC_INTERFACE"
export OC_USER="$OC_USER"
export OC_SSO_ARGS_DEFAULT="$OC_SSO_ARGS_DEFAULT"
export OC_SSO_ARGS="${OC_SSO_ARGS:-$OC_SSO_ARGS_DEFAULT}"
export OC_AUTHGROUP="${OC_AUTHGROUP:-}"
export OC_USERAGENT="${OC_USERAGENT:-}"
export OC_EXTRA_ARGS="${OC_EXTRA_ARGS:-}"
export OC_TOTP_SECRET="$OC_TOTP_SECRET"
export DISPLAY=":1"
ENVFILE

  # Build openconnect command
  OPENCONNECT_CMD="/usr/sbin/openconnect --protocol=anyconnect --servercert $OC_SERVERCERT --interface $OC_INTERFACE --script /usr/share/vpnc-scripts/vpnc-script"
  [[ -n "${OC_AUTHGROUP:-}" ]] && OPENCONNECT_CMD+=" --authgroup $OC_AUTHGROUP"
  [[ -n "${OC_USERAGENT:-}" ]] && OPENCONNECT_CMD+=" --useragent $OC_USERAGENT"
  [[ -n "${OC_EXTRA_ARGS:-}" ]] && OPENCONNECT_CMD+=" ${OC_EXTRA_ARGS}"
  echo "export OPENCONNECT_CMD=\"$OPENCONNECT_CMD\"" >> /etc/vpn.env

  cat > /usr/local/bin/vpn_connect << 'VPNCMD'
#!/usr/bin/env bash
source /etc/vpn.env
echo "[$(date)] Starting VPN connection..."

# openconnect-sso reads TOTP from keyring automatically
if [[ -n "$OC_USER" ]]; then
  echo "" | openconnect-sso -s "$OC_URL" ${OC_SSO_ARGS:-$OC_SSO_ARGS_DEFAULT} -- $OPENCONNECT_CMD
else
  openconnect-sso -s "$OC_URL" ${OC_SSO_ARGS:-$OC_SSO_ARGS_DEFAULT} -- $OPENCONNECT_CMD
fi
VPNCMD
  chmod +x /usr/local/bin/vpn_connect
}

# Create VPN runner script that keeps shell open
create_vpn_script() {
  cat > /tmp/vpn-runner.sh << 'VPNSCRIPT'
#!/usr/bin/env bash
cd /root

echo "============================================"
echo "  Cistech VPN Container"
echo "============================================"
echo ""
echo "Commands:"
echo "  vpn_connect  - Start/restart VPN connection"
echo "  Ctrl+C       - Stop auto-reconnect and drop to shell"
echo ""
echo "Starting VPN with auto-reconnect..."
echo ""

while true; do
  vpn_connect
  echo ""
  echo "[$(date)] VPN disconnected. Reconnecting in 10 seconds..."
  echo "(Press Ctrl+C to stop auto-reconnect)"
  sleep 10
done
VPNSCRIPT
  chmod +x /tmp/vpn-runner.sh
}

start_gui() {
  mkdir -p /root/.vnc
  x11vnc -storepasswd "$VNC_PASSWORD" /root/.vnc/pass >/dev/null 2>&1 || true
  rm -f /tmp/.X1-lock /tmp/.X11-unix/X1 2>/dev/null || true
  Xvfb "$DISPLAY_ADDR" -screen 0 ${XVFB_WxHxD:-1280x800x24} +extension RANDR &
  pids+=($!)
  sleep 0.5
  export DISPLAY="$DISPLAY_ADDR"
  fluxbox >/tmp/fluxbox.log 2>&1 &
  pids+=($!)
  x11vnc -display "$DISPLAY_ADDR" -rfbauth /root/.vnc/pass -forever -shared -rfbport 5900 -quiet &
  pids+=($!)
  websockify --web=/usr/share/novnc/ 0.0.0.0:"$NOVNC_PORT" localhost:5900 >/tmp/websockify.log 2>&1 &
  pids+=($!)
}

start_vpn_terminal() {
  # Start xterm with VPN script
  sleep 1
  xterm -fa 'Monospace' -fs 11 -bg black -fg white -geometry 120x35+50+50 \
    -T "Cistech VPN" -e /tmp/vpn-runner.sh &
  pids+=($!)
}

start_cloudflared() {
  case "$CLOUDFLARED_MODE" in
    token)
      [ -n "$CLOUDFLARED_TOKEN" ] && cloudflared tunnel run --token "$CLOUDFLARED_TOKEN" >/tmp/cloudflared.log 2>&1 &
      pids+=($!)
      ;;
    config)
      cloudflared --no-autoupdate --config /etc/cloudflared/config.yml tunnel run >/tmp/cloudflared.log 2>&1 &
      pids+=($!)
      ;;
    off|*)
      ;;
  esac
}

start_ssh_tunnel() {
  [ "$SSH_TUNNEL_ENABLE" = "1" ] || return 0
  IFS=',' read -ra LINES <<< "$SSH_FORWARDS"
  args=(-N -o StrictHostKeyChecking=no -o ServerAliveInterval=60)
  for m in "${LINES[@]}"; do args+=(-L "$m"); done
  ssh "${args[@]}" "$SSH_DEST" &
  pids+=($!)
}

setup_nat() {
  (
    for i in {1..60}; do
      if ip link show "$OC_INTERFACE" >/dev/null 2>&1; then
        sysctl -w net.ipv4.ip_forward=1 >/dev/null
        iptables -t nat -C POSTROUTING -o "$OC_INTERFACE" -j MASQUERADE 2>/dev/null || \
          iptables -t nat -A POSTROUTING -o "$OC_INTERFACE" -j MASQUERADE
        echo "NAT enabled on $OC_INTERFACE"

        # Trigger host routing service
        if [ -d /runtime ]; then
          touch /runtime/restart-routing
          echo "Host routing trigger sent"
        fi
        break
      fi
      sleep 2
    done
  ) &
}

trap 'kill 0' INT TERM

# Always start GUI now
setup_keyring
create_vpn_command
create_vpn_script
start_gui
start_vpn_terminal
setup_nat
start_cloudflared
start_ssh_tunnel

wait