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

# --- locate conf next to the script -----------------------------------------------------------
BASE_DIR="$(cd "$(dirname "$0")" && pwd)"
CONF_FILE="${CONF_FILE:-"$BASE_DIR/$(basename "$0" .sh).conf"}"
[[ -f "$CONF_FILE" ]] || { echo "❌ Config not found: $CONF_FILE" >&2; exit 2; }
# shellcheck source=/dev/null
source "$CONF_FILE"

# --- defaults ---------------------------------------------------------------------------------
CADDYFILE="${CADDYFILE:-/etc/caddy/Caddyfile}"
CADDY_BIN="${CADDY_BIN:-caddy}"
JQ_BIN="${JQ_BIN:-jq}"

CONNECT_TIMEOUT="${CONNECT_TIMEOUT:-7}"
MAX_TIME="${MAX_TIME:-15}"
ACCEPT_STATUS_RE='^2|3'
USER_AGENT="${USER_AGENT:-CaddyHealthCheck/2.0}"
CURL_BIN="${CURL_BIN:-curl}"

KUMA_PUSH_URL="${KUMA_PUSH_URL:-}"
KUMA_ONLY_DOWN="${KUMA_ONLY_DOWN:-0}"

INCLUDE_RE="${INCLUDE_RE:-}"
EXCLUDE_RE="${EXCLUDE_RE:-}"
EXTRA_SITES_FILE="${EXTRA_SITES_FILE:-}"

# Resiliency/overrides
FORCE_HTTP1_RE="${FORCE_HTTP1_RE:-}"
INSECURE_TLS_RE="${INSECURE_TLS_RE:-}"
HEALTH_PATHS="${HEALTH_PATHS:-}"
DIAG="${DIAG:-0}"

# --- helpers ----------------------------------------------------------------------------------
say(){ printf '%s\n' "$*"; }

url_encode() {
  local s="$1" i c out=
  for (( i=0; i<${#s}; i++ )); do
    c="${s:$i:1}"
    case "$c" in
      [a-zA-Z0-9.~_-]) out+="$c" ;;
      ' ') out+='%20' ;;
      *) printf -v out '%s%%%02X' "$out" "'$c" ;;
    esac
  done
  printf '%s' "$out"
}

valid_host_token() {
  local t="$1"
  [[ -n "$t" ]] || return 1
  [[ "$t" == *'*'* || "$t" == *'{'* || "$t" == *'}'* || "$t" == *'$'* ]] && return 1
  [[ "$t" == "("* || "$t" == ")"* || "$t" == "}" || "$t" == "{" ]] && return 1
  [[ "$t" =~ ^:?[0-9]+$ ]] && return 1
  return 0
}

to_url() {
  local host="$1"
  # if host already has scheme
  if [[ "$host" =~ ^https?:// ]]; then
    printf '%s\n' "$host"
  elif [[ "$host" =~ :80$ ]]; then
    printf 'http://%s\n' "${host%:80}"
  elif [[ "$host" =~ :443$ ]]; then
    printf 'https://%s\n' "${host%:443}"
  else
    printf 'https://%s\n' "$host"
  fi
}

valid_url() {
  local u="$1"
  [[ "$u" =~ ^https?://[^[:space:]/]+(/.*)?$ ]] || return 1
  local host="${u#*://}"; host="${host%%/*}"
  [[ "$host" != *'*'* && "$host" != *'{'* && "$host" != *'}'* && "$host" != *'$'* ]] || return 1
  [[ -n "$host" && "$host" != "}" ]] || return 1
  return 0
}

# --- extraction via Caddy JSON (preferred) ----------------------------------------------------
extract_sites_via_caddy_adapt() {
  command -v "$CADDY_BIN" >/dev/null 2>&1 || return 1
  command -v "$JQ_BIN"    >/dev/null 2>&1 || return 1

  # Generate JSON from your Caddyfile (imports & snippets resolved)
  # Note: we capture stdout; errors (if any) go to stderr and cause non-zero exit.
  local json
  if ! json="$("$CADDY_BIN" adapt --config "$CADDYFILE" --adapter caddyfile 2>/dev/null)"; then
    return 1
  fi

  # Pull every host that appears anywhere:
  # - Any matcher object containing "host": ["a", "b", ...]
  # - TLS automation policy subjects (certificate subjects)
  # We flatten, dedupe, drop wildcards, and print one per line.
  echo "$json" | "$JQ_BIN" -r '
    def hosts_in(obj):
      (obj | .. | objects | .match? // empty)               # collect any .match arrays
      | .[]? | .host? // empty | .[]?;                      # pull .host lists
    def tls_subjects(obj):
      (obj.apps.tls.automation.policies[]?.subjects[]?) // empty;

    [
      hosts_in(.) ,
      tls_subjects(.)
    ]
    | map(select(. != null))
    | map(select(. | type == "string"))
    | map(select(test("^[*]") | not))   # drop wildcards like *.example.com
    | unique
    | .[]
  ' 2>/dev/null
}

# --- legacy Caddyfile parser (fallback) -------------------------------------------------------
extract_sites_via_awk() {
  awk '
    /^[[:space:]]*$/ {next}
    /^[[:space:]]*#/ {next}
    /^[[:space:]]*\{/ {next}
    /^[[:space:]]*}/ {next}
    /^[[:space:]]*\(/ {next}
    /^[[:space:]]*import[[:space:]]/ {next}
    /^@/ {next}
    /^[^[:space:]]/ {
      line=$0
      sub(/\{.*$/,"",line)
      gsub(/[[:space:]]+/," ",line)
      gsub(/,[[:space:]]*/,",",line)
      gsub(/[[:space:]]*,/,",",line)
      gsub(/[[:space:]]+$/,"",line)
      gsub(/^[[:space:]]+/,"",line)
      if (line!="") print line
    }
  ' "$CADDYFILE" 2>/dev/null

  # Also parse @name host ... lines
  awk '
    /^[[:space:]]*#/ {next}
    /^[[:space:]]*$/ {next}
    /^[[:space:]]*@[^[:space:]]+[[:space:]]+host[[:space:]]+/ {
      line=$0
      sub(/[[:space:]]+#.*$/,"",line)
      sub(/^[[:space:]]*@[^[:space:]]+[[:space:]]+host[[:space:]]+/,"",line)
      gsub(/[[:space:]]+/," ",line)
      gsub(/[[:space:]]+$/,"",line)
      gsub(/^[[:space:]]+/,"",line)
      if (line!="") print line
    }
  ' "$CADDYFILE" 2>/dev/null
}

load_extra_sites() {
  [[ -n "${EXTRA_SITES_FILE:-}" && -f "$EXTRA_SITES_FILE" ]] || return 0
  grep -vE '^\s*(#|$)' "$EXTRA_SITES_FILE" | awk '{$1=$1;print}' | sort -u
}

with_health_path() {
  local url="$1"
  local host="${url#*://}"; host="${host%%/*}"
  local path=""

  # 1) Per-host override
  for kv in $HEALTH_PATHS; do
    local k="${kv%%=*}" v="${kv#*=}"
    if [[ "$k" == "$host" ]]; then path="$v"; break; fi
  done

  # 2) Global default if none matched
  if [[ -z "$path" && -n "${DEFAULT_HEALTH_PATH:-}" ]]; then
    path="$DEFAULT_HEALTH_PATH"
  fi

  # 3) Return URL
  if [[ -n "$path" ]]; then
    printf '%s://%s%s\n' "${url%%://*}" "$host" "$path"
  else
    printf '%s\n' "$url"
  fi
}

needs_http1()   { local u="$1"; [[ -n "$FORCE_HTTP1_RE"   && "$u" =~ $FORCE_HTTP1_RE   ]]; }
needs_insecure(){ local u="$1"; [[ -n "$INSECURE_TLS_RE" && "$u" =~ $INSECURE_TLS_RE ]]; }

diag_hint() {
  [[ "$DIAG" = "1" ]] || return 0
  local url="$1"
  local host="${url#*://}"; host="${host%%/*}"
  if command -v getent >/dev/null 2>&1; then
    getent ahosts "$host" >/dev/null || echo "   • DNS: no A/AAAA found for $host"
  fi
  if command -v timeout >/dev/null 2>&1; then
    timeout 3 bash -c "exec 3<>/dev/tcp/$host/443" 2>/dev/null \
      || echo "   • TCP: cannot open 443 to $host"
  fi
}

probe_url() {
  local url_raw="$1"
  local url; url="$(with_health_path "$url_raw")"

  # Build curl flags
  local curl_flags=(-A "$USER_AGENT" -L -sS
                  --max-time "$MAX_TIME" --connect-timeout "$CONNECT_TIMEOUT")
  # add health header(s)
  if [[ -n "${EXTRA_HEADERS:-}" ]]; then
    for h in $EXTRA_HEADERS; do curl_flags+=(-H "$h"); done
  fi

  # temp files
  local body_file; body_file="$(mktemp)"
  local code_file; code_file="$(mktemp)"
  local http_code

  # fetch (capture body + code)
  if ! "$CURL_BIN" "${curl_flags[@]}" -D /dev/null -o "$body_file" -w "%{http_code}" "$url" >"$code_file" 2>/tmp/curl_err.$$; then
    http_code="000"
  else
    http_code="$(cat "$code_file" 2>/dev/null || echo 000)"
  fi

  # Retry with HTTP/1.1 on TLS/ALPN flakiness or if code looks bogus
  if [[ "$http_code" == "000" ]] || grep -qiE 'tls|ssl|alpn|protocol error|handshake' /tmp/curl_err.$$ 2>/dev/null; then
    if "$CURL_BIN" --http1.1 "${curl_flags[@]}" -D /dev/null -o "$body_file" -w "%{http_code}" "$url" >"$code_file" 2>>/tmp/curl_err.$$; then
      http_code="$(cat "$code_file" 2>/dev/null || echo 000)"
    fi
  fi

  rm -f "$code_file" /tmp/curl_err.$$

  # Numeric success check: ONLY 200–399 is OK
  if [[ "$http_code" =~ ^[0-9]{3}$ ]]; then
    if (( http_code >= 200 && http_code < 400 )); then
      # Optional: fail if body contains known error markers even with 200/3xx
      if [[ -n "${ALSO_FAIL_IF_BODY_MATCH:-}" ]] && grep -qiE "$ALSO_FAIL_IF_BODY_MATCH" "$body_file"; then
        echo "BAD:$http_code"
        rm -f "$body_file"
        return 1
      fi
      echo "OK:$http_code"
      rm -f "$body_file"
      return 0
    fi
  fi

  echo "BAD:$http_code"
  rm -f "$body_file"
  return 1
}

kuma_push() {
  local status="$1"  # up|down
  local msg="$2"
  [[ -z "$KUMA_PUSH_URL" ]] && return 0
  local enc_msg; enc_msg="$(url_encode "$msg")"
  "$CURL_BIN" -fsS --max-time 10 "$KUMA_PUSH_URL?status=$status&msg=$enc_msg" >/dev/null || true
}

# --- MAIN -------------------------------------------------------------------------------------
[[ -r "$CADDYFILE" ]] || { echo "❌ Cannot read Caddyfile: $CADDYFILE" >&2; exit 2; }

declare -a RAW_HOSTS=()
if mapfile -t _out < <(extract_sites_via_caddy_adapt); then
  RAW_HOSTS+=("${_out[@]}")
else
  mapfile -t _out < <(extract_sites_via_awk)
  RAW_HOSTS+=("${_out[@]}")
fi

# plus extra sites
if EXTRA=$(load_extra_sites); then
  while IFS= read -r x; do RAW_HOSTS+=("$x"); done <<<"$EXTRA"
fi

# Normalize → URLs, filter, validate, dedupe
declare -A seen=()
declare -a SITES=()
for token in "${RAW_HOSTS[@]}"; do
  for piece in $token; do
    valid_host_token "$piece" || continue
    url="$(to_url "$piece")"
    valid_url "$url" || continue
    [[ -n "$INCLUDE_RE" && ! "$url" =~ $INCLUDE_RE ]] && continue
    [[ -n "$EXCLUDE_RE" &&  "$url" =~ $EXCLUDE_RE ]] && continue
    if [[ -z "${seen[$url]:-}" ]]; then
      seen[$url]=1
      SITES+=("$url")
    fi
  done
done

if [[ ${#SITES[@]} -eq 0 ]]; then
  say "ℹ️ No valid sites detected after parsing/filters."
  [[ "$KUMA_ONLY_DOWN" = "1" ]] || kuma_push up "No sites to probe"
  exit 0
fi

FAILS=()
say "🔎 Probing ${#SITES[@]} site(s) from $CADDYFILE ..."
for url in "${SITES[@]}"; do
  if probe_url "$url"; then
    say "✅ OK  $url"
  else
    say "❌ BAD $url"
    FAILS+=("$url")
  fi
done

if [[ ${#FAILS[@]} -eq 0 ]]; then
  say "🎉 All sites healthy."
  [[ "$KUMA_ONLY_DOWN" = "1" ]] || kuma_push up "All ${#SITES[@]} sites healthy"
  exit 0
else
  msg="Failures: ${#FAILS[@]}/${#SITES[@]} → $(printf '%s, ' "${FAILS[@]}" | sed 's/, $//')"
  say "🚨 $msg"
  kuma_push down "$msg"
  exit 1
fi