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

# -----------------------------------------------------------------------------
# cotg_curl_stats.ansiblelike.sh
#
# Stress-test helper that mimics the Ansible playbook logic:
#   - "uri" to fetch the APK index with retries/delay/until
#   - optional fallback "uri" to a WP JSON endpoint (best-effort)
#   - (optional) "get_url"-like APK download with retries, without *.apk.1 files
#
# Positional args (kept compatible with your original script):
#   1: N          (default: 300)
#   2: SLEEP_SEC  (default: 0.5)
#
# Environment variables:
#   UA=ansible-httpget
#   RETRIES=5
#   DELAY=5
#   JSON_URL=...                # like code_fetch_apk_url_json (optional)
#   DOWNLOAD_BASE=...           # like code_download_url (optional, used if DO_DOWNLOAD=1)
#   DO_DOWNLOAD=0|1             # if 1, download the latest armv8a APK each loop
#   DOWNLOAD_DIR=...            # where to store downloaded APKs (default: tmpdir/apk)
#   DOWNLOAD_TIMEOUT=60         # curl max-time
#   CURL_HTTP=--http1.1         # force HTTP version (Ansible's urllib is typically HTTP/1.1)
#   CURL_EXTRA_ARGS=""          # extra curl args (space-separated)
# -----------------------------------------------------------------------------

URL="https://www.appdevforall.org/codeonthego/"
N="${1:-500}"
SLEEP_SEC="${2:-0.5}"

UA="${UA:-ansible-httpget}"
RETRIES="${RETRIES:-5}"
DELAY="${DELAY:-5}"
JSON_URL="${JSON_URL:-https://www.appdevforall.org/wp-json/wp/v2/pages/2223}"
DOWNLOAD_BASE="${DOWNLOAD_BASE:-}"
DO_DOWNLOAD="${DO_DOWNLOAD:-0}"
DOWNLOAD_TIMEOUT="${DOWNLOAD_TIMEOUT:-60}"
CURL_HTTP="${CURL_HTTP:---http1.1}"
CURL_EXTRA_ARGS="${CURL_EXTRA_ARGS:-}"

# Patterns (defaults aligned to the playbook vars)
APK_RE="${APK_RE:-CodeOnTheGo-.*?armv8a\\.apk}"
CF_RE='(cloudflare|cf-chl|just a moment|attention required)'

# "uri" logic in your playbook treats these as terminal statuses for the index fetch
ALLOWED_INDEX_STATUS=(200 403 404)

tmpdir="$(mktemp -d)"
trap 'rm -rf "$tmpdir"' EXIT

DOWNLOAD_DIR="${DOWNLOAD_DIR:-$tmpdir/apk}"
mkdir -p "$DOWNLOAD_DIR"

# Stats
declare -A status_count=()
declare -A apk_count=()
declare -A cf_count=()
declare -A blocked_count=()
declare -A retry_hist_index=()
declare -A retry_hist_json=()
declare -A retry_hist_apk=()

min_size=""
max_size=""
sum_size=0

note_size() {
  local sz="$1"
  [[ -z "$min_size" || "$sz" -lt "$min_size" ]] && min_size="$sz"
  [[ -z "$max_size" || "$sz" -gt "$max_size" ]] && max_size="$sz"
  sum_size=$((sum_size + sz))
}

in_list() {
  local x="$1"; shift
  local v
  for v in "$@"; do
    [[ "$x" == "$v" ]] && return 0
  done
  return 1
}

# curl -> files, return status code (0 means "undefined" like Ansible when status missing)
# Writes headers/body even if status is 0.
curl_fetch() {
  local url="$1" body="$2" hdr="$3" accept_header="${4:-}"
  local rc status
  local -a extra

  # Allow tuning curl behavior to get closer to what Ansible's urllib does.
  # (Example: force HTTP/1.1 instead of opportunistic HTTP/2.)
  extra=("$CURL_HTTP")
  if [[ -n "$CURL_EXTRA_ARGS" ]]; then
    # shellcheck disable=SC2206
    extra+=( $CURL_EXTRA_ARGS )
  fi

  : >"$hdr"
  : >"$body"

  # Disable -e locally to collect rc + still return 0 status on failures
  set +e
  if [[ -n "$accept_header" ]]; then
    curl -sS -L -A "$UA" -m "$DOWNLOAD_TIMEOUT" "${extra[@]}" -H "Accept: $accept_header" -D "$hdr" -o "$body" "$url"
  else
    curl -sS -L -A "$UA" -m "$DOWNLOAD_TIMEOUT" "${extra[@]}" -D "$hdr" -o "$body" "$url"
  fi
  rc=$?
  set -e

  # status = last HTTP status in the chain
  status="$(awk '/^HTTP\//{code=$2} END{print code+0}' "$hdr" 2>/dev/null)"

  # If curl failed hard, treat as "status undefined"
  if [[ $rc -ne 0 || -z "$status" || "$status" -eq 0 ]]; then
    echo 0
    return 0
  fi

  echo "$status"
}

# Mimic:
#   retries: RETRIES
#   delay:   DELAY
#   until:   status is defined and status in [200,403,404]
#   failed_when: status is not defined or status not in [200,403,404]
# Returns: "<status> <attempts>" where attempts is how many tries were used.
fetch_index_like_ansible() {
  local url="$1" body="$2" hdr="$3"
  local attempt status

  for ((attempt=1; attempt<=RETRIES; attempt++)); do
    status="$(curl_fetch "$url" "$body" "$hdr")"

    # until: status defined AND status in allowed
    if [[ "$status" -ne 0 ]] && in_list "$status" "${ALLOWED_INDEX_STATUS[@]}"; then
      echo "$status $attempt"
      return 0
    fi

    [[ $attempt -lt $RETRIES ]] && sleep "$DELAY"
  done

  # exhausted retries; status may be 0 or non-allowed
  echo "$status $RETRIES"
}

# Mimic your JSON fallback:
#   retries: RETRIES
#   delay:   DELAY
#   until:   status is defined
#   failed_when: false   (best-effort)
# Returns: "<status> <attempts>" with status 0 if never got any HTTP status.
fetch_json_best_effort() {
  local url="$1" body="$2" hdr="$3"
  local attempt status

  for ((attempt=1; attempt<=RETRIES; attempt++)); do
    status="$(curl_fetch "$url" "$body" "$hdr" 'application/json')"

    if [[ "$status" -ne 0 ]]; then
      echo "$status $attempt"
      return 0
    fi

    [[ $attempt -lt $RETRIES ]] && sleep "$DELAY"
  done

  echo "0 $RETRIES"
}

# Extract content.rendered from a WP JSON response (similar to Ansible's parsed json)
# Prints rendered HTML to stdout, or nothing on failure.
extract_wp_rendered() {
  local json_file="$1"

  if command -v python3 >/dev/null 2>&1; then
    python3 - "$json_file" <<'PY'
import json, sys
p = sys.argv[1]
try:
    with open(p, 'r', encoding='utf-8', errors='replace') as f:
        data = json.load(f)
    rendered = data.get('content', {}).get('rendered', '')
    if isinstance(rendered, str):
        sys.stdout.write(rendered)
except Exception:
    pass
PY
  else
    # Very rough fallback (not JSON-safe); prefer python3.
    sed -n 's/.*"rendered"[[:space:]]*:[[:space:]]*"\(.*\)".*/\1/p' "$json_file" | head -n 1
  fi
}

# get_url-like download:
#   retries: RETRIES
#   delay:   DELAY
#   until:   succeeded
# Ensures final file name is exactly dest (downloads to temp then atomic mv).
# Returns: "<ok(0|1)> <attempts> <http_status>"
download_apk_like_ansible() {
  local url="$1" dest="$2"
  local attempt status tmp_part hdr

  hdr="$tmpdir/apk.hdr"

  for ((attempt=1; attempt<=RETRIES; attempt++)); do
    tmp_part="${dest}.part.$$"

    status="$(curl_fetch "$url" "$tmp_part" "$hdr")"

    if [[ "$status" -eq 200 && -s "$tmp_part" ]]; then
      mv -f "$tmp_part" "$dest"
      echo "1 $attempt $status"
      return 0
    fi

    rm -f "$tmp_part"
    [[ $attempt -lt $RETRIES ]] && sleep "$DELAY"
  done

  echo "0 $RETRIES ${status:-0}"
}

echo "URL: $URL"
echo "Iterations: $N   Sleep: ${SLEEP_SEC}s   UA: $UA"
echo "RETRIES: $RETRIES   DELAY: ${DELAY}s"
[[ -n "$JSON_URL" ]] && echo "JSON_URL: $JSON_URL"
[[ -n "$DOWNLOAD_BASE" ]] && echo "DOWNLOAD_BASE: $DOWNLOAD_BASE"
[[ "$DO_DOWNLOAD" == "1" ]] && echo "DO_DOWNLOAD: 1 (dir: $DOWNLOAD_DIR)"
echo

for ((i=1; i<=N; i++)); do
  body="$tmpdir/body.$i.html"
  hdr="$tmpdir/hdr.$i.txt"
  json_body="$tmpdir/json.$i.json"
  json_hdr="$tmpdir/jsonhdr.$i.txt"
  rendered_body="$tmpdir/rendered.$i.html"

  # 1) Fetch index like Ansible 'uri' task
  read -r idx_status idx_attempts < <(fetch_index_like_ansible "$URL" "$body" "$hdr")
  idx_retries_used=$((idx_attempts - 1))

  # classify status for stats
  if [[ "$idx_status" -eq 0 ]]; then
    status_key="undefined"   # like "status is undefined"
  else
    status_key="$idx_status"
  fi
  status_count["$status_key"]=$(( ${status_count["$status_key"]:-0} + 1 ))
  retry_hist_index["$idx_retries_used"]=$(( ${retry_hist_index["$idx_retries_used"]:-0} + 1 ))

  size="$(wc -c < "$body" | tr -d ' ')"
  note_size "$size"

  # 2) Detect whether APK links exist
  has_apk=0
  if grep -Eqo "$APK_RE" "$body"; then
    has_apk=1
  fi

  # 3) Optional JSON fallback if no APK link found
  json_status=0
  json_attempts=0
  if [[ "$has_apk" -eq 0 && -n "$JSON_URL" ]]; then
    read -r json_status json_attempts < <(fetch_json_best_effort "$JSON_URL" "$json_body" "$json_hdr")
    json_retries_used=$((json_attempts - 1))
    retry_hist_json["$json_retries_used"]=$(( ${retry_hist_json["$json_retries_used"]:-0} + 1 ))

    # If 200 and we can extract rendered HTML, replace the body (mimic set_fact update)
    if [[ "$json_status" -eq 200 ]]; then
      extracted="$(extract_wp_rendered "$json_body" || true)"
      if [[ -n "$extracted" ]]; then
        printf '%s' "$extracted" > "$rendered_body"
        body_to_parse="$rendered_body"
      else
        body_to_parse="$body"
      fi
    else
      body_to_parse="$body"
    fi

    # re-check APK links after fallback
    if grep -Eqo "$APK_RE" "$body_to_parse"; then
      has_apk=1
    fi
  else
    body_to_parse="$body"
  fi

  # 4) Mimic code_blocked_by_cdn condition:
  #    blocked if 403 OR (apk_missing even after fallback)
  blocked=0
  if [[ "$idx_status" -eq 403 || "$has_apk" -eq 0 ]]; then
    blocked=1
  fi
  blocked_count["$blocked"]=$(( ${blocked_count["$blocked"]:-0} + 1 ))

  # Additional signal (not in Ansible, but useful)
  if grep -Eqi "$CF_RE" "$body_to_parse"; then
    cf_count["cf_like"]=$(( ${cf_count["cf_like"]:-0} + 1 ))
  else
    cf_count["no_cf"]=$(( ${cf_count["no_cf"]:-0} + 1 ))
  fi

  if [[ "$has_apk" -eq 1 ]]; then
    apk_count["apk_found"]=$(( ${apk_count["apk_found"]:-0} + 1 ))
  else
    apk_count["apk_missing"]=$(( ${apk_count["apk_missing"]:-0} + 1 ))
  fi

  # 5) Optional APK download like Ansible get_url
  dl_ok=0
  dl_attempts=0
  dl_status=0
  if [[ "$DO_DOWNLOAD" == "1" && "$blocked" -eq 0 ]]; then
    # Match playbook behavior: Jinja2 `sort` is lexicographic.
    apk_name="$(grep -Eo "$APK_RE" "$body_to_parse" | LC_ALL=C sort | tail -n 1 || true)"
    if [[ -n "$apk_name" ]]; then
      # Match Ansible playbook behavior: download comes from a separate base domain.
      # If DOWNLOAD_BASE is not set, fall back to the index URL's directory.
      if [[ -n "$DOWNLOAD_BASE" ]]; then
        apk_url="${DOWNLOAD_BASE%/}/$apk_name"
      else
        apk_url="${URL%/}/$apk_name"
      fi
      apk_dest="$DOWNLOAD_DIR/$apk_name"
      read -r dl_ok dl_attempts dl_status < <(download_apk_like_ansible "$apk_url" "$apk_dest")
      dl_retries_used=$((dl_attempts - 1))
      retry_hist_apk["$dl_retries_used"]=$(( ${retry_hist_apk["$dl_retries_used"]:-0} + 1 ))
    fi
  fi

  # progress line
  printf "[%03d/%03d] idx=%s(attempts=%s) apk=%s blocked=%s json=%s(attempts=%s)\r" \
    "$i" "$N" \
    "${idx_status:-0}" "$idx_attempts" \
    "$([[ "$has_apk" -eq 1 ]] && echo Y || echo N)" \
    "$blocked" \
    "${json_status:-0}" "${json_attempts:-0}"

  sleep "$SLEEP_SEC"
done

echo -e "\n\n==== SUMMARY (Ansible-like) ===="

echo "Index status counts (0 means status undefined):"
for k in "${!status_count[@]}"; do
  printf "  %-10s %d\n" "$k" "${status_count[$k]}"
done | sort -k2,2nr

echo
echo "APK armv8a links:"
for k in "${!apk_count[@]}"; do
  printf "  %-12s %d\n" "$k" "${apk_count[$k]}"
done | sort -k2,2nr

echo
echo "Blocked (Ansible condition: idx==403 OR apk_missing):"
printf "  blocked=1  %d\n" "${blocked_count[1]:-0}"
printf "  blocked=0  %d\n" "${blocked_count[0]:-0}"

echo
echo "Retry histogram (retries used):"
echo "  Index fetch (retries used -> count):"
for k in "${!retry_hist_index[@]}"; do
  printf "    %s -> %d\n" "$k" "${retry_hist_index[$k]}"
done | sort -n -k1,1

if [[ -n "$JSON_URL" ]]; then
  echo "  JSON fallback (retries used -> count):"
  for k in "${!retry_hist_json[@]}"; do
    printf "    %s -> %d\n" "$k" "${retry_hist_json[$k]}"
  done | sort -n -k1,1
fi

if [[ "$DO_DOWNLOAD" == "1" ]]; then
  echo "  APK download (retries used -> count):"
  for k in "${!retry_hist_apk[@]}"; do
    printf "    %s -> %d\n" "$k" "${retry_hist_apk[$k]}"
  done | sort -n -k1,1
fi

echo
avg_size=$((sum_size / N))
echo "Body size bytes: min=$min_size  max=$max_size  avg=$avg_size"

echo
if [[ "$DO_DOWNLOAD" == "1" ]]; then
  echo "Downloaded APKs are in: $DOWNLOAD_DIR"
fi

echo "Note: files are saved under $tmpdir and will be auto-removed at exit (trap)."