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

# Usage:
#   host="100.123.68.125"
#   remote_addr="127.0.0.1"   # or "0.0.0.0" or any ip bound on remote
#   port=8000
#   kill_remote_vscode "$host" "$remote_addr" "$port"
#
# Returns:
#   0 = killed successfully
#   1 = nothing to kill / no listener found
#   2 = listener found but not Code/code-tunnel (skipped)
#   3 = error (ssh failure / parsing failure)
kill_remote_vscode_windows() {
  local HOST="$1" REMOTE_ADDR="$2" PORT="$3"
  local netstat_cmd="netstat -ano | findstr \"${REMOTE_ADDR}:${PORT}\" | findstr LISTENING || true"
  local netline
  netline=$(ssh "$HOST" "$netstat_cmd") || { echo "ERROR: ssh netstat failed" >&2; return 3; }
  if [[ -z "$netline" ]]; then
    echo "No listener found on ${REMOTE_ADDR}:${PORT} on ${HOST}."; return 1; fi
  local pid
  pid=$(printf '%s\n' "$netline" | awk '{print $NF}' | head -n1)
  if [[ -z "$pid" ]]; then
    echo "Failed to parse PID from netstat output:" >&2; printf '%s\n' "$netline" >&2; return 3; fi
  local taskline tasklist_cmd="tasklist /FI \"PID eq ${pid}\" /FO CSV /NH"
  taskline=$(ssh "$HOST" "$tasklist_cmd") || { echo "ERROR: ssh tasklist failed" >&2; return 3; }
  if printf '%s\n' "$taskline" | grep -qi "No tasks are running"; then echo "PID $pid not present."; return 1; fi
  local procname lower
  procname=$(printf '%s\n' "$taskline" | awk -F',' '{gsub(/"/,"",$1); print $1}' | head -n1)
  lower=$(printf '%s' "$procname" | tr '[:upper:]' '[:lower:]')
  if [[ "$lower" == "code.exe" || "$lower" == "code" || "$lower" == "code-tunnel.exe" || "$lower" == "code-tunnel" ]]; then
    echo "Killing windows PID $pid on $HOST ..."
    ssh "$HOST" "taskkill /PID ${pid} /F" && { echo "Killed PID $pid."; return 0; }
    echo "Failed to kill PID $pid." >&2; return 3
  else
    echo "PID $pid belongs to '$procname' — not Code/code-tunnel. Skipping kill."; return 2
  fi
}

kill_remote_vscode_unix() {
  # Linux / macOS: find PIDs listening on REMOTE_ADDR:PORT, prefer lsof then ss, validate by command line, then kill
  local HOST="$1" REMOTE_ADDR="$2" PORT="$3"
  # Build remote script safely using a heredoc with placeholder substitution to avoid local shell quoting issues
  local REMOTE_SCRIPT
  REMOTE_SCRIPT=$(cat <<'EOS'
set -euo pipefail
found_pids=""
if command -v lsof >/dev/null 2>&1; then
  found_pids=$(lsof -nP -iTCP:__PORT__ -sTCP:LISTEN -t 2>/dev/null || true)
fi
if [ -z "$found_pids" ] && command -v ss >/dev/null 2>&1; then
  found_pids=$(ss -ltnp 2>/dev/null | awk -v port=__PORT__ '$4 ~ ":"port"$" {print}' | sed -n 's/.*pid=\([0-9]\+\).*/\1/p')
fi
if [ -z "$found_pids" ] && command -v pgrep >/dev/null 2>&1; then
  found_pids=$(pgrep -f "serve-web.*(--port[= ]__PORT__|:__PORT__)" || true)
fi
killed_any=1
if [ -n "$found_pids" ]; then
  for pid in $found_pids; do
    cmdline=$(ps -p "$pid" -o args= 2>/dev/null || true)
    echo "$cmdline" | grep -qiE "(code(|-tunnel).*(serve-web)|serve-web)" || continue
    echo "Killing PID $pid ($cmdline)"
    kill -TERM "$pid" 2>/dev/null || true
  done
  sleep 0.5
  for pid in $found_pids; do
    if kill -0 "$pid" 2>/dev/null; then
      kill -KILL "$pid" 2>/dev/null || true
    fi
  done
  killed_any=0
fi
exit "$killed_any"
EOS
)
  REMOTE_SCRIPT=${REMOTE_SCRIPT//__PORT__/$PORT}

  if ssh "$HOST" bash -lc "$REMOTE_SCRIPT"; then
    echo "Remote VS Code server cleaned up (Linux/macOS)."; return 0
  else
    echo "No matching VS Code server found to clean on $HOST (Linux/macOS)."; return 1
  fi
}


# ----------------------------------
# Defaults
# ----------------------------------
PORT=8000
REMOTE_ADDR="127.0.0.1"
REMOTE_PORT=8000
BIND_HOST_IPV4=false
KILL_MODE=false

# ----------------------------------
# Usage
# ----------------------------------
usage() {
  cat <<EOF
Usage: $0 [options] <host>

Options:
  -p, --port <port>         Local/forward port (default: 8000)
  -r, --remote <addr[:p]>   Remote bind address and optional port (default: 127.0.0.1)
  -b, --bindhost            Use detected IPv4 of <host> as remote bind address
                            (mutually exclusive with --remote)
  -k, --kill                Kill the VS Code server on the given port/host
  -h, --help                Show this help

Examples:
  $0 myhost
  $0 -p 8080 myhost
  $0 -r 0.0.0.0 myhost
  $0 -r 192.168.1.10:9000 myhost
  $0 --kill -r 0.0.0.0:9000 myhost
  $0 -b myhost              # bind the VS Code server to the host's IPv4
EOF
  exit 1
}

# ----------------------------------
# Parse args
# ----------------------------------
REMOTE_SPEC=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    -p|--port)
      shift; PORT="${1:-}"; [[ -z "$PORT" ]] && usage ;;
    -r|--remote)
      shift; REMOTE_SPEC="${1:-}"; [[ -z "$REMOTE_SPEC" ]] && usage ;;
    -b|--bindhost)
      BIND_HOST_IPV4=true ;;
    -k|--kill)
      KILL_MODE=true ;;
    -h|--help)
      usage ;;
    -*)
      echo "Unknown option: $1" >&2; usage ;;
    *)
      HOST="$1" ;;
  esac
  shift
done

[[ -z "${HOST:-}" ]] && usage

# ----------------------------------
# Validate mutual exclusivity
# ----------------------------------
if [[ "$BIND_HOST_IPV4" == true && -n "$REMOTE_SPEC" ]]; then
  echo "Error: --bindhost (-b) cannot be used together with --remote (-r)." >&2
  exit 1
fi

# ----------------------------------
# Resolve host IPv4
# ----------------------------------
IPV4=$(ping -c 1 "$HOST" 2>/dev/null | grep -E -o "\(.*\):" | grep -E -o "[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+") || true
if [[ -n "$IPV4" ]]; then
  echo "Host $HOST is mapped to $IPV4"
else
  IPV4="127.0.0.1"
  echo "Failed to resolve IPv4 for $HOST; using 127.0.0.1"
fi

# ----------------------------------
# Parse remote spec (ip[:port])
# ----------------------------------
if [[ -n "$REMOTE_SPEC" ]]; then
  if [[ "$REMOTE_SPEC" == *:* ]]; then
    REMOTE_ADDR="${REMOTE_SPEC%%:*}"
    REMOTE_PORT="${REMOTE_SPEC##*:}"
  elif [[ "$REMOTE_SPEC" =~ ^[0-9]+$ ]]; then
    REMOTE_PORT="$REMOTE_SPEC"
  else
    REMOTE_ADDR="$REMOTE_SPEC"
  fi
fi

# If user wants to bind to detected IPv4
if [[ "$BIND_HOST_IPV4" == true ]]; then
  REMOTE_ADDR="$IPV4"
fi

# ----------------------------------
# Detect remote OS
# ----------------------------------
echo "Detecting OS on $HOST ..."
REMOTE_OS=$(ssh -o BatchMode=yes -o ConnectTimeout=5 "$HOST" 'uname' | head -n 1)

case "$REMOTE_OS" in
  *[Ll]inux*)
    OS_TYPE="linux"
    CODE_CMD="code"
    ;;
  *[Dd]arwin*)
    OS_TYPE="mac"
    CODE_CMD="/Applications/Visual\\ Studio\\ Code.app/Contents/Resources/app/bin/code"
    ;;
  *[Mm][Ss][Yy][Ss]* | *[Mm][Ii][Nn][Gg][Ww]* | *[Cc][Yy][Gg][Ww]* | *[Ww]indows*)
    OS_TYPE="windows"
    CODE_CMD="\"d:/App/Microsoft VS Code/bin/code\""
    ;;
  *)
    OS_TYPE="unknown"
    CODE_CMD="code"
    ;;
esac

echo "Detected remote OS: $OS_TYPE ($REMOTE_OS)"

cleanup_remote() {
  echo "Cleaning up remote VS Code server on $HOST (host: $REMOTE_ADDR, port: $REMOTE_PORT)..."
  case "$OS_TYPE" in
    windows)
      kill_remote_vscode_windows "$HOST" "$REMOTE_ADDR" "$REMOTE_PORT" || true ;;
    linux|mac)
      kill_remote_vscode_unix "$HOST" "$REMOTE_ADDR" "$REMOTE_PORT" || true ;;
    *)
      echo "Unknown remote OS; attempted generic cleanup..." ;;
  esac
}

# Ensure we attempt cleanup on Ctrl+C or script exit
trap cleanup_remote INT TERM EXIT

# ----------------------------------
# Kill mode (explicit command)
# ----------------------------------
if [[ "$KILL_MODE" == true ]]; then
  cleanup_remote
  exit 0
fi


echo -e "\n------------------------------------------------------------------\n"


# ----------------------------------
# Start VS Code server
# ----------------------------------
echo "Starting VS Code server on $HOST (bind ${REMOTE_ADDR}:${REMOTE_PORT})..."
ssh -L "localhost:${PORT}:${REMOTE_ADDR}:${REMOTE_PORT}" "$HOST" \
  "${CODE_CMD} serve-web --without-connection-token --accept-server-license-terms --host ${REMOTE_ADDR} --port ${REMOTE_PORT}"
