#!/usr/bin/env bash
# Guard 110 Prometheus alert rules against stale deploys.
#
# This script is intentionally narrow: it only restores the canonical alert
# rules file when required recovery/backup rules disappear from live Prometheus
# or when the active file differs from the canonical copy.

set -uo pipefail

HOST_LABEL="${HOST_LABEL:-110}"
PROMETHEUS_URL="${PROMETHEUS_URL:-http://127.0.0.1:9090}"
CURRENT_RULES="${CURRENT_RULES:-/home/wooo/monitoring/alerts.yml}"
CANONICAL_RULES="${CANONICAL_RULES:-/home/wooo/monitoring/alerts-unified.canonical.yml}"
TEXTFILE="${TEXTFILE:-/home/wooo/node_exporter_textfiles/prometheus_rule_drift_guard.prom}"
LOG_FILE="${LOG_FILE:-/home/wooo/logs/prometheus-rule-drift-guard.log}"
PROMETHEUS_CONTAINER="${PROMETHEUS_CONTAINER:-prometheus}"
PROMETHEUS_RUNTIME_RULES="${PROMETHEUS_RUNTIME_RULES:-/etc/prometheus/alerts.yml}"
PROMETHEUS_COMPOSE_FILE="${PROMETHEUS_COMPOSE_FILE:-/home/wooo/monitoring/docker-compose.yml}"
PROMETHEUS_COMPOSE_SERVICE="${PROMETHEUS_COMPOSE_SERVICE:-prometheus}"
ALLOW_RUNTIME_RECREATE="${ALLOW_RUNTIME_RECREATE:-1}"

REQUIRED_RULES=(
  "BackupCredentialEscrowEvidenceMissing"
  "BackupExpectedJobMissing"
  "GiteaRestoreDrillMissingOrFailed"
  "awoooi_recovery_core_ready"
  "awoooi_recovery_dr_offsite_ready"
  "ColdStartRecoveryBlocked"
)

log() {
  mkdir -p "$(dirname "$LOG_FILE")" 2>/dev/null || true
  printf '[%s] %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*" >>"$LOG_FILE"
}

write_textfile() {
  local status="$1"
  local repaired="$2"
  local missing_count="$3"
  local matches_canonical="$4"
  local runtime_matches_current="$5"
  local tmp
  mkdir -p "$(dirname "$TEXTFILE")" 2>/dev/null || true
  tmp="$(mktemp "${TEXTFILE}.tmp.XXXXXX")" || return 0
  cat >"$tmp" <<EOF
# HELP awoooi_prometheus_rule_drift_guard_last_run_timestamp Unix timestamp of the last Prometheus rule drift guard run.
# TYPE awoooi_prometheus_rule_drift_guard_last_run_timestamp gauge
awoooi_prometheus_rule_drift_guard_last_run_timestamp{host="${HOST_LABEL}",status="${status}"} $(date +%s)
# HELP awoooi_prometheus_rule_drift_guard_repaired Whether the guard restored canonical Prometheus rules on the last run.
# TYPE awoooi_prometheus_rule_drift_guard_repaired gauge
awoooi_prometheus_rule_drift_guard_repaired{host="${HOST_LABEL}"} ${repaired}
# HELP awoooi_prometheus_rule_drift_guard_missing_required_count Number of required live rules missing after the last check.
# TYPE awoooi_prometheus_rule_drift_guard_missing_required_count gauge
awoooi_prometheus_rule_drift_guard_missing_required_count{host="${HOST_LABEL}"} ${missing_count}
# HELP awoooi_prometheus_rule_drift_guard_current_matches_canonical Whether active alerts.yml matches canonical copy.
# TYPE awoooi_prometheus_rule_drift_guard_current_matches_canonical gauge
awoooi_prometheus_rule_drift_guard_current_matches_canonical{host="${HOST_LABEL}"} ${matches_canonical}
# HELP awoooi_prometheus_rule_drift_guard_runtime_matches_current Whether the Prometheus container mount matches the active host rules file.
# TYPE awoooi_prometheus_rule_drift_guard_runtime_matches_current gauge
awoooi_prometheus_rule_drift_guard_runtime_matches_current{host="${HOST_LABEL}"} ${runtime_matches_current}
EOF
  chmod 0644 "$tmp" 2>/dev/null || true
  mv "$tmp" "$TEXTFILE" 2>/dev/null || rm -f "$tmp"
}

rules_missing_count() {
  python3 - "$PROMETHEUS_URL" "${REQUIRED_RULES[@]}" <<'PY'
import json
import sys
import urllib.request

base_url = sys.argv[1].rstrip("/")
required = set(sys.argv[2:])
try:
    with urllib.request.urlopen(f"{base_url}/api/v1/rules", timeout=8) as response:
        payload = json.loads(response.read().decode("utf-8"))
    if payload.get("status") != "success":
        raise RuntimeError(payload)
    loaded = {
        str(rule.get("name") or rule.get("alert") or rule.get("record")): rule
        for group in payload.get("data", {}).get("groups") or []
        for rule in group.get("rules") or []
    }
    missing = required - set(loaded)
    unhealthy = {name for name in required & set(loaded) if loaded[name].get("health") != "ok"}
    print(len(missing | unhealthy))
except Exception as exc:
    print(f"QUERY_FAILED:{exc}")
PY
}

wait_for_required_rules() {
  local attempt missing
  for attempt in $(seq 1 30); do
    missing="$(rules_missing_count)"
    if [ "$missing" = "0" ]; then
      log "runtime rules converged after attempt=${attempt}"
      return 0
    fi
    sleep 2
  done
  log "runtime rules did not converge within 60s: ${missing}"
  return 1
}

matches_canonical() {
  if [ ! -f "$CURRENT_RULES" ] || [ ! -f "$CANONICAL_RULES" ]; then
    echo 0
    return
  fi
  if cmp -s "$CURRENT_RULES" "$CANONICAL_RULES"; then
    echo 1
  else
    echo 0
  fi
}

file_sha256() {
  sha256sum "$1" 2>/dev/null | awk '{print $1}'
}

runtime_rules_sha256() {
  docker exec "${PROMETHEUS_CONTAINER}" sha256sum "${PROMETHEUS_RUNTIME_RULES}" 2>/dev/null | awk '{print $1}'
}

runtime_matches_current() {
  local current_hash runtime_hash
  current_hash="$(file_sha256 "${CURRENT_RULES}")"
  runtime_hash="$(runtime_rules_sha256)"
  if [ -n "${current_hash}" ] && [ "${current_hash}" = "${runtime_hash}" ]; then
    echo 1
  else
    echo 0
  fi
}

validate_canonical_rules() {
  local candidate_path rc
  candidate_path="/tmp/awoooi-prometheus-rule-guard-candidate.$$.yml"
  docker cp "${CANONICAL_RULES}" "${PROMETHEUS_CONTAINER}:${candidate_path}" >/dev/null 2>&1 || return 1
  docker exec "${PROMETHEUS_CONTAINER}" promtool check rules "${candidate_path}" >>"${LOG_FILE}" 2>&1
  rc=$?
  docker exec -u 0 "${PROMETHEUS_CONTAINER}" rm -f "${candidate_path}" >/dev/null 2>&1 || true
  return "${rc}"
}

recreate_prometheus_runtime() {
  [ "${ALLOW_RUNTIME_RECREATE}" = "1" ] || return 1
  [ -f "${PROMETHEUS_COMPOSE_FILE}" ] || return 1
  docker compose -f "${PROMETHEUS_COMPOSE_FILE}" up -d --no-deps --force-recreate \
    "${PROMETHEUS_COMPOSE_SERVICE}" >>"${LOG_FILE}" 2>&1
}

restore_rules() {
  local backup_path needs_recreate
  backup_path="${CURRENT_RULES}.guard.bak.$(date +%Y%m%d%H%M%S)"
  if ! cp "$CURRENT_RULES" "$backup_path" 2>/dev/null; then
    log "failed to create active rules rollback copy"
    return 1
  fi
  if ! validate_canonical_rules; then
    log "canonical promtool validation failed"
    return 1
  fi

  if ! cat "$CANONICAL_RULES" >"$CURRENT_RULES"; then
    log "failed to update active rules file"
    return 1
  fi

  needs_recreate=0
  [ "$(runtime_matches_current)" -eq 1 ] || needs_recreate=1
  if [ "${needs_recreate}" -eq 1 ]; then
    log "runtime bind mount differs from active file; recreating Prometheus only"
    if ! recreate_prometheus_runtime; then
      cat "$backup_path" >"$CURRENT_RULES" 2>/dev/null || true
      log "Prometheus recreate failed; active file rolled back"
      return 1
    fi
  elif ! curl -fsS -X POST "${PROMETHEUS_URL}/-/reload" >/dev/null; then
    cat "$backup_path" >"$CURRENT_RULES" 2>/dev/null || true
    curl -fsS -X POST "${PROMETHEUS_URL}/-/reload" >/dev/null 2>&1 || true
    log "Prometheus reload failed; active file rolled back"
    return 1
  fi

  if wait_for_required_rules && [ "$(runtime_matches_current)" -eq 1 ]; then
    return 0
  fi

  cat "$backup_path" >"$CURRENT_RULES" 2>/dev/null || true
  if [ "${needs_recreate}" -eq 1 ]; then
    recreate_prometheus_runtime >/dev/null 2>&1 || true
  else
    curl -fsS -X POST "${PROMETHEUS_URL}/-/reload" >/dev/null 2>&1 || true
  fi
  log "post-apply verifier failed; active rules rolled back"
  return 1
}

main() {
  if [ ! -f "$CANONICAL_RULES" ]; then
    log "canonical rules missing: ${CANONICAL_RULES}"
    write_textfile "canonical_missing" 0 999 0 0
    return 1
  fi

  local missing before_matches before_runtime_matches repaired after_missing after_matches after_runtime_matches
  missing="$(rules_missing_count)"
  before_matches="$(matches_canonical)"
  before_runtime_matches="$(runtime_matches_current)"
  repaired=0

  if [[ "$missing" == QUERY_FAILED:* ]]; then
    log "Prometheus query failed: ${missing}"
    write_textfile "query_failed" 0 999 "$before_matches" "$before_runtime_matches"
    return 1
  fi

  if [ "$missing" -gt 0 ] || [ "$before_matches" -eq 0 ] || [ "$before_runtime_matches" -eq 0 ]; then
    log "rule drift detected: missing=${missing} current_matches_canonical=${before_matches} runtime_matches_current=${before_runtime_matches}; restoring"
    if restore_rules; then
      repaired=1
    else
      log "restore failed"
      write_textfile "restore_failed" 0 "$missing" "$before_matches" "$before_runtime_matches"
      return 1
    fi
  fi

  after_missing="$(rules_missing_count)"
  after_matches="$(matches_canonical)"
  after_runtime_matches="$(runtime_matches_current)"
  if [[ "$after_missing" == QUERY_FAILED:* ]]; then
    log "post-restore Prometheus query failed: ${after_missing}"
    write_textfile "post_query_failed" "$repaired" 999 "$after_matches" "$after_runtime_matches"
    return 1
  fi

  if [ "$after_missing" -eq 0 ] && [ "$after_matches" -eq 1 ] && [ "$after_runtime_matches" -eq 1 ]; then
    write_textfile "ok" "$repaired" "$after_missing" "$after_matches" "$after_runtime_matches"
    log "ok repaired=${repaired}"
    return 0
  fi

  log "still drifted after check: missing=${after_missing} current_matches_canonical=${after_matches} runtime_matches_current=${after_runtime_matches}"
  write_textfile "drifted" "$repaired" "$after_missing" "$after_matches" "$after_runtime_matches"
  return 1
}

main "$@"