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

usage() {
  cat <<'USAGE'
Usage:
  bash scripts/load-speed.sh 'https://rustpad.example.com/p/test-note-features'
  bash scripts/load-speed.sh 'https://rustpad.example.com/w/team'
  bash scripts/load-speed.sh 'https://rustpad.example.com/w/team/n/note'

The script will prompt for the RustPad account login and password.
If the pad/workspace also has resource password protection, it will ask for that password separately.
USAGE
}

for cmd in curl jq wc mktemp; do
  if ! command -v "$cmd" >/dev/null 2>&1; then
    echo "ERROR: required command not found: $cmd" >&2
    exit 1
  fi
done

if [[ $# -ne 1 ]]; then
  usage
  exit 1
fi

INPUT_URL="$1"

if [[ ! "$INPUT_URL" =~ ^(https?://[^/]+)(/.*)?$ ]]; then
  echo "ERROR: pass a full URL starting with http:// or https://" >&2
  exit 1
fi

BASE="${BASH_REMATCH[1]}"
PATH_WITH_QUERY="${BASH_REMATCH[2]:-/}"
PATH_ONLY="${PATH_WITH_QUERY%%\#*}"
PATH_ONLY="${PATH_ONLY%%\?*}"
[[ "$PATH_ONLY" != "/" ]] && PATH_ONLY="${PATH_ONLY%/}"

TYPE=""
RESOURCE_KIND=""
RESOURCE_SLUG=""
API_BASE=""
WORKSPACE_SLUG=""
NOTE_SLUG=""
PAD_SLUG=""

if [[ "$PATH_ONLY" =~ ^/p/([^/]+)$ ]]; then
  TYPE="pad"
  RESOURCE_KIND="pad"
  PAD_SLUG="${BASH_REMATCH[1]}"
  RESOURCE_SLUG="$PAD_SLUG"
  API_BASE="$BASE/api/pads/$PAD_SLUG"
elif [[ "$PATH_ONLY" =~ ^/w/([^/]+)/n/([^/]+)$ ]]; then
  TYPE="workspace-note"
  RESOURCE_KIND="workspace"
  WORKSPACE_SLUG="${BASH_REMATCH[1]}"
  NOTE_SLUG="${BASH_REMATCH[2]}"
  RESOURCE_SLUG="$WORKSPACE_SLUG"
  API_BASE="$BASE/api/workspaces/$WORKSPACE_SLUG/notes/$NOTE_SLUG"
elif [[ "$PATH_ONLY" =~ ^/w/([^/]+)$ ]]; then
  TYPE="workspace"
  RESOURCE_KIND="workspace"
  WORKSPACE_SLUG="${BASH_REMATCH[1]}"
  RESOURCE_SLUG="$WORKSPACE_SLUG"
  API_BASE="$BASE/api/workspaces/$WORKSPACE_SLUG"
else
  echo "ERROR: unsupported RustPad URL: $PATH_ONLY" >&2
  echo "Supported: /p/<pad>, /w/<workspace>, /w/<workspace>/n/<note>" >&2
  exit 1
fi

TMP_DIR="$(mktemp -d)"
COOKIE_JAR="$TMP_DIR/cookies.txt"
CSRF_BODY="$TMP_DIR/csrf.json"
LOGIN_BODY="$TMP_DIR/login.json"
INFO_BODY="$TMP_DIR/info.json"
UNLOCK_BODY="$TMP_DIR/unlock.json"

cleanup() {
  rm -rf "$TMP_DIR"
}
trap cleanup EXIT INT TERM

touch "$COOKIE_JAR"
chmod 600 "$COOKIE_JAR"

CURL_COMMON=(
  --silent
  --show-error
  --compressed
  --connect-timeout 10
  --max-time 120
  -b "$COOKIE_JAR"
  -c "$COOKIE_JAR"
)

human_bytes() {
  local bytes="$1"
  if (( bytes >= 1048576 )); then
    awk -v b="$bytes" 'BEGIN { printf "%.2f MiB", b / 1048576 }'
  elif (( bytes >= 1024 )); then
    awk -v b="$bytes" 'BEGIN { printf "%.1f KiB", b / 1024 }'
  else
    printf '%s B' "$bytes"
  fi
}

measure_to_file() {
  local name="$1"
  local outfile="$2"
  shift 2

  local stats code transfer ttfb total body_bytes
  stats="$(curl "${CURL_COMMON[@]}" "$@" \
    -o "$outfile" \
    -w '%{http_code}|%{size_download}|%{time_starttransfer}|%{time_total}')"

  IFS='|' read -r code transfer ttfb total <<< "$stats"
  body_bytes="$(wc -c < "$outfile" | tr -d '[:space:]')"

  printf '%-20s HTTP=%-3s transfer=%10s body=%10s TTFB=%8ss total=%8ss\n' \
    "$name" "$code" "$(human_bytes "${transfer%.*}")" "$(human_bytes "$body_bytes")" "$ttfb" "$total"

  LAST_HTTP="$code"
}

measure() {
  local name="$1"
  shift
  local outfile="$TMP_DIR/measure-${name//[^a-zA-Z0-9_-]/_}.body"
  measure_to_file "$name" "$outfile" "$@"
}

api_post_json() {
  local url="$1"
  local json="$2"
  local outfile="$3"
  local code

  code="$(curl "${CURL_COMMON[@]}" \
    -H 'Content-Type: application/json' \
    -H "x-rustpad-csrf: $CSRF" \
    --data "$json" \
    -o "$outfile" \
    -w '%{http_code}' \
    "$url")"
  LAST_HTTP="$code"
}

fetch_csrf() {
  local code
  code="$(curl "${CURL_COMMON[@]}" \
    -o "$CSRF_BODY" \
    -w '%{http_code}' \
    "$BASE/api/security/csrf")"

  if [[ "$code" != "200" ]]; then
    echo "ERROR: CSRF endpoint returned HTTP $code" >&2
    cat "$CSRF_BODY" >&2 || true
    exit 1
  fi

  CSRF="$(jq -r '.token // empty' "$CSRF_BODY")"
  if [[ -z "$CSRF" ]]; then
    echo "ERROR: server did not return a CSRF token" >&2
    exit 1
  fi
}

login_account() {
  local login password payload code message nickname

  printf 'Login/e-mail: '
  IFS= read -r login
  if [[ -z "$login" ]]; then
    echo "ERROR: login cannot be empty" >&2
    exit 1
  fi

  printf 'Password: '
  IFS= read -rs password
  printf '\n'
  if [[ -z "$password" ]]; then
    echo "ERROR: password cannot be empty" >&2
    exit 1
  fi

  payload="$(jq -nc --arg email "$login" --arg password "$password" '{email:$email,password:$password}')"
  unset password

  code="$(curl "${CURL_COMMON[@]}" \
    -H 'Content-Type: application/json' \
    -H "x-rustpad-csrf: $CSRF" \
    --data "$payload" \
    -o "$LOGIN_BODY" \
    -w '%{http_code}' \
    "$BASE/api/auth/login")"
  unset payload

  if [[ "$code" != "200" ]]; then
    message="$(jq -r '.error // empty' "$LOGIN_BODY" 2>/dev/null || true)"
    echo "ERROR: login failed (HTTP $code)${message:+: $message}" >&2
    exit 1
  fi

  nickname="$(jq -r '.nickname // empty' "$LOGIN_BODY")"
  echo "Logged in${nickname:+ as $nickname}."
}

unlock_resource_if_needed() {
  local info_file="$1"
  local protected access_level resource_password payload code message

  protected="$(jq -r '.protected // false' "$info_file" 2>/dev/null || echo false)"
  access_level="$(jq -r '.access_level // empty' "$info_file" 2>/dev/null || true)"

  if [[ "$protected" != "true" || "$access_level" != "none" ]]; then
    return 0
  fi

  printf 'Resource password (%s): ' "$RESOURCE_KIND"
  IFS= read -rs resource_password
  printf '\n'
  if [[ -z "$resource_password" ]]; then
    echo "ERROR: this resource requires a password" >&2
    exit 1
  fi

  payload="$(jq -nc \
    --arg kind "$RESOURCE_KIND" \
    --arg slug "$RESOURCE_SLUG" \
    --arg password "$resource_password" \
    '{kind:$kind,slug:$slug,password:$password}')"
  unset resource_password

  api_post_json "$BASE/api/access-token" "$payload" "$UNLOCK_BODY"
  code="$LAST_HTTP"
  unset payload

  if [[ "$code" != "200" ]]; then
    message="$(jq -r '.error // empty' "$UNLOCK_BODY" 2>/dev/null || true)"
    echo "ERROR: resource unlock failed (HTTP $code)${message:+: $message}" >&2
    exit 1
  fi

  echo "Resource unlocked."
}

print_header() {
  echo
  echo "Target: $INPUT_URL"
  echo "Type:   $TYPE"
  echo "Base:   $BASE"
  echo
}

fetch_csrf
login_account
print_header

echo '=== PAGE / SESSION ==='
measure "page_html" "$BASE$PATH_WITH_QUERY"
measure "auth_me" "$BASE/api/auth/me"

echo
case "$TYPE" in
  pad)
    echo '=== INITIAL NOTE LOAD ==='
    measure_to_file "note_info" "$INFO_BODY" "$API_BASE"
    if [[ "$LAST_HTTP" != "200" ]]; then
      echo "ERROR: note_info failed; cannot continue reliably." >&2
      jq -r '.error // empty' "$INFO_BODY" 2>/dev/null >&2 || true
      exit 1
    fi

    unlock_resource_if_needed "$INFO_BODY"

    measure "editor_color" "$API_BASE/editor-color"
    measure "favorite_status" "$BASE/api/auth/favorites/status?kind=pad&slug=$PAD_SLUG"
    measure "files" \
      -X PUT \
      -H 'Content-Type: application/json' \
      -H "x-rustpad-csrf: $CSRF" \
      --data '{"access_token":null}' \
      "$API_BASE/files"

    echo
    echo '=== ON DEMAND ==='
    measure "history" \
      -X POST \
      -H 'Content-Type: application/json' \
      -H "x-rustpad-csrf: $CSRF" \
      --data '{"access_token":null}' \
      "$API_BASE/history"
    ;;

  workspace-note)
    echo '=== INITIAL NOTE LOAD ==='
    measure_to_file "note_info" "$INFO_BODY" "$API_BASE"
    if [[ "$LAST_HTTP" != "200" ]]; then
      echo "ERROR: note_info failed; cannot continue reliably." >&2
      jq -r '.error // empty' "$INFO_BODY" 2>/dev/null >&2 || true
      exit 1
    fi

    unlock_resource_if_needed "$INFO_BODY"

    measure "editor_color" "$API_BASE/editor-color"
    measure "favorite_status" "$BASE/api/auth/favorites/status?kind=note&slug=$NOTE_SLUG&workspace_slug=$WORKSPACE_SLUG"
    measure "files" \
      -X PUT \
      -H 'Content-Type: application/json' \
      -H "x-rustpad-csrf: $CSRF" \
      --data '{"access_token":null}' \
      "$API_BASE/files"

    echo
    echo '=== ON DEMAND ==='
    measure "history" \
      -X POST \
      -H 'Content-Type: application/json' \
      -H "x-rustpad-csrf: $CSRF" \
      --data '{"access_token":null}' \
      "$API_BASE/history"
    ;;

  workspace)
    echo '=== INITIAL WORKSPACE LOAD ==='
    measure_to_file "workspace_info" "$INFO_BODY" "$API_BASE"
    if [[ "$LAST_HTTP" != "200" ]]; then
      echo "ERROR: workspace_info failed; cannot continue reliably." >&2
      jq -r '.error // empty' "$INFO_BODY" 2>/dev/null >&2 || true
      exit 1
    fi

    unlock_resource_if_needed "$INFO_BODY"

    measure "workspace_open" \
      -X POST \
      -H 'Content-Type: application/json' \
      -H "x-rustpad-csrf: $CSRF" \
      --data '{"access_token":null}' \
      "$API_BASE/open?q=&page=1&per_page=25"
    ;;
esac

echo
echo 'Note: WebSocket payloads are not included in these curl measurements.'