#!/bin/bash
# Pi-Kit post-prep smoke test (HTTP/HTTPS/API/firstboot/services)

set -euo pipefail

PIKIT_HOST="${PIKIT_HOST:-pikit.local}"
PIKIT_USER="${PIKIT_USER:-dietpi}"
PIKIT_SSH_KEY="${PIKIT_SSH_KEY:-$HOME/.ssh/pikit}"
PIKIT_SSH_OPTS="${PIKIT_SSH_OPTS:-}"
PIKIT_HTTP_URL="${PIKIT_HTTP_URL:-http://$PIKIT_HOST}"
PIKIT_HTTPS_URL="${PIKIT_HTTPS_URL:-https://$PIKIT_HOST}"
PIKIT_API_URL="${PIKIT_API_URL:-http://127.0.0.1:4000}"

LOCAL_ONLY=0
ERRORS=0
WARNINGS=0
REMOTE_MODE=0

usage() {
  cat <<'USAGE'
Usage: pikit-smoke-test.sh [--local]

Runs a quick post-prep smoke test:
- HTTP/HTTPS reachable
- API reachable and returns JSON
- firstboot state done
- core services active (nginx, pikit-api, dietpi-dashboard-frontend)

Options:
  --local   Run locally on the Pi (skip SSH)
  --help    Show this help

Env:
  PIKIT_HOST, PIKIT_USER, PIKIT_SSH_KEY, PIKIT_SSH_OPTS
  PIKIT_HTTP_URL, PIKIT_HTTPS_URL
USAGE
}

status() {
  local level="$1"
  shift
  printf '[%s] %s\n' "$level" "$*"
  case "$level" in
    FAIL) ERRORS=$((ERRORS + 1)) ;;
    WARN) WARNINGS=$((WARNINGS + 1)) ;;
  esac
}

section() {
  printf '\n== %s ==\n' "$1"
}

is_dietpi() {
  grep -qi "dietpi" /etc/os-release 2>/dev/null
}

parse_args() {
  for arg in "$@"; do
    case "$arg" in
      --local) LOCAL_ONLY=1 ;;
      --help|-h) usage; exit 0 ;;
      *)
        echo "[FAIL] Unknown argument: $arg" >&2
        usage
        exit 1
        ;;
    esac
  done
}

remote_cmd() {
  local cmd="$1"
  if [ "$LOCAL_ONLY" -eq 1 ] || is_dietpi; then
    bash -c "$cmd"
  else
    ssh -i "$PIKIT_SSH_KEY" $PIKIT_SSH_OPTS -o ConnectTimeout=10 "${PIKIT_USER}@${PIKIT_HOST}" "$cmd"
  fi
}

extract_json_line() {
  awk 'BEGIN{found=0} /^[[:space:]]*[{[]/ {print; found=1; exit} END{if(!found) exit 0}'
}

json_get() {
  local key="$1"
  if command -v python3 >/dev/null 2>&1; then
    python3 -c 'import json,sys
key=sys.argv[1]
try:
    data=json.load(sys.stdin)
except Exception:
    print("")
    sys.exit(1)
val=data.get(key, "")
if isinstance(val, bool):
    print("true" if val else "false")
else:
    print(val)
' "$key"
  elif command -v jq >/dev/null 2>&1; then
    jq -r --arg key "$key" '.[$key] // empty'
  else
    cat >/dev/null
    echo ""
    return 1
  fi
}

check_http() {
  local url="$1"
  local label="$2"
  if curl -fsS --max-time 5 "$url" >/dev/null; then
    status OK "$label reachable"
  else
    status FAIL "$label not reachable"
  fi
}

check_https() {
  local url="$1"
  local label="$2"
  if curl -kfsS --max-time 5 "$url" >/dev/null; then
    status OK "$label reachable"
  else
    status FAIL "$label not reachable"
  fi
}

check_api() {
  local url="$1"
  local body
  if ! body="$(remote_cmd "curl -fsS --max-time 5 $url")"; then
    status FAIL "API not reachable: $url"
    return
  fi
  if [ "$REMOTE_MODE" -eq 1 ]; then
    body="$(printf "%s\n" "$body" | extract_json_line)"
  fi
  if [ -z "$body" ]; then
    status FAIL "API response empty or not JSON"
    return
  fi
  if command -v python3 >/dev/null 2>&1; then
    if printf "%s" "$body" | python3 -c 'import json,sys
try:
    data=json.load(sys.stdin)
except Exception:
    sys.exit(1)
for key in ("services","hostname","uptime_seconds"):
    if key in data:
        sys.exit(0)
sys.exit(1)
'
    then
      status OK "API responds with JSON"
    else
      status WARN "API response did not include expected fields"
    fi
  else
    status WARN "python3 missing; API JSON check skipped"
  fi
}

check_firstboot() {
  local url="$1"
  local body state error_present
  if ! body="$(remote_cmd "curl -fsS --max-time 5 $url")"; then
    status FAIL "firstboot API not reachable"
    return
  fi
  if [ "$REMOTE_MODE" -eq 1 ]; then
    body="$(printf "%s\n" "$body" | extract_json_line)"
  fi
  if [ -z "$body" ]; then
    status FAIL "firstboot status invalid or missing"
    return
  fi
  state="$(printf "%s" "$body" | json_get "state" || true)"
  error_present="$(printf "%s" "$body" | json_get "error_present" || true)"
  if [ -z "$state" ]; then
    status FAIL "firstboot status invalid or missing"
    return
  fi
  if [ "$state" = "done" ] && [ "$error_present" != "true" ]; then
    status OK "firstboot completed"
  else
    status FAIL "firstboot not complete (state=$state error=$error_present)"
  fi
}

check_services() {
  local services=("nginx" "pikit-api" "dietpi-dashboard-frontend")
  for svc in "${services[@]}"; do
    if remote_cmd "systemctl is-active --quiet $svc"; then
      status OK "$svc active"
    else
      status FAIL "$svc not active"
    fi
  done
}

check_ports() {
  local cmd="ss -lnt | awk '{print \$4}' | grep -E ':(80|443|5252|5253)\$' | sort -u"
  local out
  if out="$(remote_cmd "$cmd" 2>/dev/null)"; then
    if echo "$out" | grep -q ":80" && echo "$out" | grep -q ":443"; then
      status OK "ports 80/443 listening"
    else
      status WARN "ports 80/443 not both listening"
    fi
  else
    status WARN "unable to check ports"
  fi
}

finalize() {
  section "Summary"
  status OK "warnings: $WARNINGS"
  status OK "errors: $ERRORS"
  if [ "$ERRORS" -gt 0 ]; then
    echo "[FAIL] Smoke test failed."
    exit 1
  fi
  echo "[OK] Smoke test passed."
}

main() {
  parse_args "$@"
  if [ "$LOCAL_ONLY" -eq 0 ] && ! is_dietpi; then
    REMOTE_MODE=1
  fi

  section "HTTP/HTTPS"
  check_http "$PIKIT_HTTP_URL" "HTTP"
  check_https "$PIKIT_HTTPS_URL" "HTTPS"

  section "API"
  check_api "$PIKIT_API_URL/api/status"

  section "Firstboot"
  check_firstboot "$PIKIT_API_URL/api/firstboot"

  section "Services"
  check_services

  section "Ports"
  check_ports

  finalize
}

main "$@"