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

BASE_URL="https://remotesh.polybase.app"

if [[ "${BASE_URL}" != https://* ]]; then
  echo "Refusing to run over non-HTTPS URL: ${BASE_URL}" >&2
  echo "Set ALLOW_HTTP=1 on server only for local development." >&2
  exit 1
fi

if ! command -v curl >/dev/null 2>&1; then
  echo "curl is required." >&2
  exit 1
fi

banner() {
  printf '\033[38;5;51m'
  cat <<'BANNER'
+--------------------------------------------------------------+
|   remotesh // secure relay terminal                          |
|   role: CONTROLLER                                           |
+--------------------------------------------------------------+
BANNER
  printf '\033[0m'
}

api() {
  local method="$1"
  local path="$2"
  local token="${3:-}"
  if [[ -n "$token" ]]; then
    curl -fsS -X "$method" "$BASE_URL$path" -H "Authorization: Bearer $token"
  else
    curl -fsS -X "$method" "$BASE_URL$path"
  fi
}

post_kv() {
  local path="$1"
  local token="$2"
  local data="$3"
  curl -fsS -X POST "$BASE_URL$path" \
    -H "Authorization: Bearer $token" \
    -H "Content-Type: text/plain; charset=utf-8" \
    --data-binary "$data"
}

banner
printf '\033[38;5;246m[init]\033[0m Creating one-time session...\n'
SESSION_DATA="$(api POST /api/v1/sessions)"

getv() {
  local key="$1"
  printf '%s\n' "$SESSION_DATA" | sed -n "s/^$key=//p" | head -n1
}

SESSION_ID="$(getv SESSION_ID)"
CONTROLLER_TOKEN="$(getv CONTROLLER_TOKEN)"
TARGET_TOKEN="$(getv TARGET_TOKEN)"
EXPIRES_AT="$(getv EXPIRES_AT)"

if [[ -z "$SESSION_ID" || -z "$CONTROLLER_TOKEN" || -z "$TARGET_TOKEN" ]]; then
  echo "Failed to create session." >&2
  exit 1
fi

TARGET_CMD="curl -fsSL $BASE_URL/target.sh | bash -s -- $SESSION_ID $TARGET_TOKEN"

printf '\n\033[38;5;46m[share]\033[0m Send this exact command to your friend:\n\n'
printf '  %s\n\n' "$TARGET_CMD"
printf '\033[38;5;220m[session]\033[0m id=%s  expires=%s\n' "$SESSION_ID" "$EXPIRES_AT"

if [[ ! -r /dev/tty ]]; then
  echo "Interactive terminal is required (missing /dev/tty)." >&2
  exit 1
fi
exec 3</dev/tty

wait_for_target() {
  local tick=0
  local frames=(
    "[..........]"
    "[#.........]"
    "[##........]"
    "[###.......]"
    "[####......]"
    "[#####.....]"
    "[######....]"
    "[#######...]"
    "[########..]"
    "[#########.]"
    "[##########]"
  )

  while true; do
    local status connected closed frame
    status="$(api GET "/api/v1/sessions/$SESSION_ID/status" "$CONTROLLER_TOKEN" || true)"
    connected="$(printf '%s\n' "$status" | sed -n 's/^TARGET_CONNECTED=//p' | head -n1)"
    closed="$(printf '%s\n' "$status" | sed -n 's/^CLOSED=//p' | head -n1)"
    if [[ "$connected" == "1" ]]; then
      printf '\r\033[38;5;46m[ready]\033[0m Target connected.%*s\n' 22 ''
      return 0
    fi
    if [[ "$closed" == "1" ]]; then
      printf '\n\033[38;5;196m[error]\033[0m Session was closed before target connected.\n'
      return 1
    fi
    frame="${frames[tick % ${#frames[@]}]}"
    printf '\r\033[38;5;220m[wait]\033[0m Waiting for target to connect %s' "$frame"
    tick=$((tick + 1))
    sleep 1
  done
}

wait_for_target || exit 1
printf '\033[38;5;246mType commands. Use :help for local shortcuts.\033[0m\n\n'

RESULT_CURSOR=0

decode_b64() {
  local in="$1"
  if base64 --decode >/dev/null 2>&1 <<<"$in"; then
    printf '%s' "$in" | base64 --decode
  else
    printf '%s' "$in" | base64 -D
  fi
}

poll_until_result() {
  local wanted_seq="$1"
  while true; do
    RESULT_DATA="$(api GET "/api/v1/sessions/$SESSION_ID/results/next?after=$RESULT_CURSOR&wait=25" "$CONTROLLER_TOKEN" || true)"
    if [[ "$RESULT_DATA" == NOOP* || -z "$RESULT_DATA" ]]; then
      continue
    fi
    local seq exit_code out_b64
    seq="$(printf '%s\n' "$RESULT_DATA" | sed -n 's/^SEQ=//p' | head -n1)"
    exit_code="$(printf '%s\n' "$RESULT_DATA" | sed -n 's/^EXIT=//p' | head -n1)"
    out_b64="$(printf '%s\n' "$RESULT_DATA" | sed -n 's/^OUT_B64=//p' | head -n1)"
    if [[ -z "$seq" ]]; then
      continue
    fi
    RESULT_CURSOR="$seq"
    local output
    output="$(decode_b64 "$out_b64" || true)"
    if [[ "$seq" == "$wanted_seq" ]]; then
      printf '\033[38;5;39m[remote:%s]\033[0m exit=%s\n' "$seq" "$exit_code"
      if [[ -n "$output" ]]; then
        printf '%s\n' "$output"
      fi
      return 0
    fi
  done
}

while true; do
  printf '\033[38;5;51mremotesh>\033[0m '
  IFS= read -r line <&3 || break
  case "$line" in
    "")
      continue
      ;;
    :help)
      cat <<'HELP'
Local commands:
  :help   Show this help
  :exit   Close session
  :status Print connection status
HELP
      continue
      ;;
    :status)
      STATUS_DATA="$(api GET "/api/v1/sessions/$SESSION_ID/status" "$CONTROLLER_TOKEN")"
      printf '%s\n' "$STATUS_DATA"
      continue
      ;;
    :exit)
      api POST "/api/v1/sessions/$SESSION_ID/close" "$CONTROLLER_TOKEN" >/dev/null || true
      printf '\033[38;5;208m[done]\033[0m Session closed.\n'
      exit 0
      ;;
  esac

  CMD_B64="$(printf '%s' "$line" | base64 | tr -d '\n')"
  RESPONSE="$(post_kv "/api/v1/sessions/$SESSION_ID/commands" "$CONTROLLER_TOKEN" "CMD_B64=$CMD_B64"$'\n')"
  SEQ="$(printf '%s\n' "$RESPONSE" | sed -n 's/^SEQ=//p' | head -n1)"
  if [[ -z "$SEQ" ]]; then
    printf '\033[38;5;196m[error]\033[0m Failed to submit command.\n'
    continue
  fi
  poll_until_result "$SEQ"
done