#!/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)
- profile-specific checks (if /etc/pikit/profile.json exists)

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
}

remote_sudo_cmd() {
  remote_cmd "sudo bash -c \"$1\""
}

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
  local done_present error_file_present log_present state_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
  done_present="$(remote_sudo_cmd "test -f /var/lib/pikit/firstboot/firstboot.done && echo yes || echo no" 2>/dev/null || true)"
  error_file_present="$(remote_sudo_cmd "test -f /var/lib/pikit/firstboot/firstboot.error && echo yes || echo no" 2>/dev/null || true)"
  log_present="$(remote_sudo_cmd "test -f /var/lib/pikit/firstboot/firstboot.log && echo yes || echo no" 2>/dev/null || true)"
  state_present="$(remote_sudo_cmd "test -f /var/lib/pikit/firstboot/state.json && echo yes || echo no" 2>/dev/null || true)"
  if [ "$state" = "done" ] && [ "$error_present" != "true" ]; then
    status OK "firstboot completed"
    return
  fi
  if [ "$state" = "error" ] || [ "$error_present" = "true" ] || [ "$error_file_present" = "yes" ]; then
    status FAIL "firstboot failed (state=$state error=$error_present)"
    return
  fi
  if [ "$done_present" = "yes" ]; then
    status FAIL "firstboot state mismatch (done file present but state=$state)"
    return
  fi
  if [ "$log_present" != "yes" ] && [ "$state_present" != "yes" ]; then
    status WARN "firstboot not started yet (image prepped?)"
  else
    status WARN "firstboot in progress (state=$state)"
  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
}

load_profile() {
  local body
  body="$(remote_cmd "python3 -c 'import json, pathlib, sys; p=pathlib.Path(\"/etc/pikit/profile.json\"); print(json.dumps(json.loads(p.read_text()))) if p.exists() else sys.exit(0)'" 2>/dev/null || true)"
  body="$(printf "%s" "$body" | extract_json_line)"
  if [ -n "$body" ]; then
    printf "%s" "$body"
    return 0
  fi
  return 1
}

json_get_list() {
  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:
    sys.exit(1)
val=data.get(key, [])
if not isinstance(val, list):
    sys.exit(0)
for item in val:
    print(json.dumps(item))
' "$key"
  elif command -v jq >/dev/null 2>&1; then
    jq -c --arg key "$key" '.[$key] // [] | .[]'
  fi
}

check_profile_firewall() {
  local firewall_enable="$1"
  local ports_csv="$2"
  local ufw_out
  if [ "$firewall_enable" != "true" ]; then
    status OK "profile firewall not required"
    return
  fi
  if ! remote_cmd "test -x /usr/sbin/ufw || test -x /usr/bin/ufw"; then
    status FAIL "ufw missing (profile expects firewall enabled)"
    return
  fi
  ufw_out="$(remote_sudo_cmd "timeout 6 /usr/sbin/ufw status 2>/dev/null || /usr/bin/ufw status 2>/dev/null || ufw status 2>/dev/null" 2>/dev/null || true)"
  if ! printf "%s" "$ufw_out" | grep -qi "Status: active"; then
    status FAIL "ufw not active"
    return
  fi
  status OK "ufw active"
  if [ -n "$ports_csv" ]; then
    IFS=',' read -r -a ports <<<"$ports_csv"
    for port in "${ports[@]}"; do
      if printf "%s" "$ufw_out" | grep -Eq "(^|[[:space:]])${port}(/|[[:space:]])"; then
        status OK "ufw allows port $port"
      else
        status WARN "ufw rule missing for port $port"
      fi
    done
  fi
}

check_profile_services() {
  local services_json="$1"
  local profile_services="$2"
  if [ -z "$profile_services" ]; then
    status OK "profile services: none"
    return
  fi
  if [ -z "$services_json" ]; then
    status FAIL "services.json missing (profile services expected)"
    return
  fi
  if command -v python3 >/dev/null 2>&1; then
    local result
    result="$(SERVICES_JSON="$services_json" PROFILE_SERVICES="$profile_services" python3 - <<'PY'
import json, os, re
services_text = os.environ.get("SERVICES_JSON", "")
profile_text = os.environ.get("PROFILE_SERVICES", "")
try:
    services_data = json.loads(services_text or "[]")
except Exception:
    services_data = []
profile_lines = [line for line in profile_text.splitlines() if line.strip()]
services = [s for s in services_data if isinstance(s, dict)]
def norm(x):
    return re.sub(r"\s+", " ", str(x or "")).strip().lower()
missing = []
for line in profile_lines:
    try:
        psvc = json.loads(line)
    except Exception:
        continue
    name = norm(psvc.get("name"))
    port = str(psvc.get("port") or "")
    path = str(psvc.get("path") or "")
    found = False
    for svc in services:
        sname = norm(svc.get("name"))
        sport = str(svc.get("port") or "")
        spath = str(svc.get("path") or "")
        if name and sname == name:
            found = True
            if port and sport != port:
                print(f"WARN: service {name} port mismatch ({sport} != {port})")
            if path and spath != path:
                print(f"WARN: service {name} path mismatch ({spath} != {path})")
            break
        if port and sport == port and path and spath == path:
            found = True
            break
    if not found:
        missing.append(name or f"port {port}")
if missing:
    print("MISSING:" + ",".join(missing))
PY
)"
    if [ -z "$result" ]; then
      status OK "profile services registered"
      return
    fi
    if echo "$result" | grep -q "^MISSING:"; then
      status FAIL "profile services missing: ${result#MISSING:}"
    fi
    echo "$result" | grep "^WARN:" | while read -r line; do
      status WARN "${line#WARN: }"
    done
  else
    status WARN "python3 missing; profile service checks skipped"
  fi
}

check_dns_stack_profile() {
  section "Profile: dns-stack"
  if remote_cmd "systemctl is-active --quiet pihole-FTL"; then
    status OK "pihole-FTL active"
  else
    status FAIL "pihole-FTL not active"
  fi
  if remote_cmd "systemctl is-active --quiet unbound"; then
    status OK "unbound active"
  else
    status FAIL "unbound not active"
  fi
  if remote_cmd "ss -lnt | grep -Eq ':53[[:space:]]'"; then
    status OK "DNS port 53 listening"
  else
    status FAIL "DNS port 53 not listening"
  fi
  if remote_cmd "ss -lnt | grep -q '127.0.0.1:5335'"; then
    status OK "Unbound 5335 listening on loopback"
  else
    status WARN "Unbound 5335 not bound to loopback"
  fi
  if remote_cmd "sudo grep -q '127.0.0.1#5335' /etc/pihole/pihole.toml"; then
    status OK "Pi-hole upstream points to Unbound"
  else
    status FAIL "Pi-hole upstream not set to 127.0.0.1#5335"
  fi
  if remote_cmd "sudo grep -q 'interface: 127.0.0.1' /etc/unbound/unbound.conf.d/dietpi.conf"; then
    status OK "Unbound listens on 127.0.0.1"
  else
    status WARN "Unbound interface not 127.0.0.1"
  fi
  if remote_cmd "sudo grep -q 'port: 5335' /etc/unbound/unbound.conf.d/dietpi.conf"; then
    status OK "Unbound port 5335 configured"
  else
    status WARN "Unbound port 5335 not configured"
  fi
  if remote_sudo_cmd "test -f /etc/pihole/tls.pem && test -f /etc/pikit/certs/pikit.local.crt"; then
    local fp_pikit fp_pihole
    fp_pikit="$(remote_sudo_cmd "openssl x509 -in /etc/pikit/certs/pikit.local.crt -noout -fingerprint -sha256 | cut -d= -f2" 2>/dev/null || true)"
    fp_pihole="$(remote_sudo_cmd "openssl x509 -in /etc/pihole/tls.pem -noout -fingerprint -sha256 | cut -d= -f2" 2>/dev/null || true)"
    if [ -n "$fp_pikit" ] && [ "$fp_pikit" = "$fp_pihole" ]; then
      status OK "Pi-hole TLS matches Pi-Kit cert"
    else
      status WARN "Pi-hole TLS does not match Pi-Kit cert"
    fi
  else
    status WARN "Pi-hole TLS bundle missing"
  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

  section "Profile"
  local profile_json
  profile_json="$(load_profile || true)"
  if [ -z "$profile_json" ]; then
    status OK "profile not present; skipping profile checks"
    finalize
    exit 0
  fi
  local profile_id firewall_enable ports_csv
  profile_id="$(printf "%s" "$profile_json" | json_get "id" || true)"
  firewall_enable="$(printf "%s" "$profile_json" | json_get "firewall_enable" || true)"
  ports_csv="$(printf "%s" "$profile_json" | python3 - <<'PY' 2>/dev/null || true
import json, sys
try:
    data=json.load(sys.stdin)
except Exception:
    sys.exit(0)
ports=data.get("firewall_ports") or []
print(",".join(str(p) for p in ports))
PY
)"
  status OK "profile detected: ${profile_id:-unknown}"
  check_profile_firewall "$firewall_enable" "$ports_csv"

  local services_json profile_services_output
  services_json="$(remote_sudo_cmd "cat /etc/pikit/services.json" 2>/dev/null || true)"
  profile_services_output="$(printf "%s" "$profile_json" | json_get_list "services" || true)"
  check_profile_services "$services_json" "$profile_services_output"

  if [ "$profile_id" = "dns-stack" ]; then
    check_dns_stack_profile
  else
    status OK "profile-specific checks skipped"
  fi

  finalize
}

main "$@"