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

# ShiftOS Update Center backend (CLI)
# Philosophy: 1) Check for Updates  2) Install (download) Updates  3) Reboot and Apply Updates

log() { echo "[shift-update-center] $*"; }
warn() { echo "[shift-update-center][WARN] $*" >&2; }
die() { echo "[shift-update-center][ERROR] $*" >&2; exit 1; }

STATE_DIR="/var/lib/shiftos-update-center"
STATE_JSON="${STATE_DIR}/state.json"
OFFLINE_FLAG="${STATE_DIR}/offline_update.flag"
DEFAULT_CONF="/etc/shiftos-update-center.conf"

mkdir -p "${STATE_DIR}"

# Defaults (can be overridden by /etc/shiftos-update-center.conf)
BATCH_DAYS=14
PKG_THRESHOLD=100
SECURITY_OVERRIDE=1     # if >=1 security patches found, recommend update sooner
CHECK_INTERVAL_HOURS=6  # informational only

load_conf() {
  if [[ -f "${DEFAULT_CONF}" ]]; then
    # shellcheck disable=SC1090
    source "${DEFAULT_CONF}" || true
  fi
}

now_epoch() { date +%s; }
iso_now() { date -Is; }

read_state() {
  if [[ -f "${STATE_JSON}" ]]; then
    cat "${STATE_JSON}"
  else
    echo "{}"
  fi
}

write_state() {
  local tmp="${STATE_JSON}.tmp"
  cat >"${tmp}"
  mv -f "${tmp}" "${STATE_JSON}"
}

get_json_int() {
  local key="$1"
  python3 - <<PY 2>/dev/null || echo 0
import json,sys
d=json.load(sys.stdin)
print(int(d.get("${key}",0) or 0))
PY
}

get_json_str() {
  local key="$1"
  python3 - <<PY 2>/dev/null || echo ""
import json,sys
d=json.load(sys.stdin)
v=d.get("${key}","")
print(v if isinstance(v,str) else "")
PY
}

ensure_root() {
  [[ "${EUID:-$(id -u)}" -eq 0 ]] || die "Run as root."
}

zypper_n() {
  # Always non-interactive with auto key import
  zypper --non-interactive --gpg-auto-import-keys "$@"
}

count_updates() {
  # Returns: pending_pkg_count
  # Uses "zypper lu" which works on TW; if it fails, returns 0 and warns.
  if ! command -v zypper >/dev/null 2>&1; then
    warn "zypper not found"
    echo 0
    return 0
  fi

  # Ensure metadata current for accurate count.
  zypper_n refresh >/dev/null || true

  # Count lines starting with "v |" (zypper lu output)
  local n
  n="$(zypper_n list-updates 2>/dev/null | awk -F'|' 'BEGIN{c=0} $1~"^v"{c++} END{print c+0}')"
  echo "${n:-0}"
}

count_security_patches() {
  # Tumbleweed generally doesn't rely on patches, but if patches exist, count security ones.
  # If not supported, return 0.
  if ! command -v zypper >/dev/null 2>&1; then
    echo 0
    return 0
  fi

  if zypper_n lp -g security >/dev/null 2>&1; then
    # lp output has table; count rows starting with "v |" similar.
    zypper_n lp -g security 2>/dev/null | awk -F'|' 'BEGIN{c=0} $1~"^v"{c++} END{print c+0}'
  else
    echo 0
  fi
}

should_recommend() {
  # Inputs: pending_count security_count last_apply_epoch
  local pending="$1" security="$2" last_apply="$3"
  local now; now="$(now_epoch)"
  local due_by_time=0
  local days_since=9999

  if [[ "${last_apply}" -gt 0 ]]; then
    days_since=$(( (now - last_apply) / 86400 ))
  fi

  if [[ "${days_since}" -ge "${BATCH_DAYS}" ]]; then
    due_by_time=1
  fi

  local due_by_count=0
  if [[ "${pending}" -ge "${PKG_THRESHOLD}" ]]; then
    due_by_count=1
  fi

  local due_by_security=0
  if [[ "${SECURITY_OVERRIDE}" -gt 0 && "${security}" -ge "${SECURITY_OVERRIDE}" ]]; then
    due_by_security=1
  fi

  if [[ "${due_by_security}" -eq 1 || "${due_by_time}" -eq 1 || "${due_by_count}" -eq 1 ]]; then
    echo 1
  else
    echo 0
  fi
}

cmd_check() {
  load_conf

  local pending security
  pending="$(count_updates)"
  security="$(count_security_patches)"

  local state last_apply
  state="$(read_state)"
  last_apply="$(printf '%s' "${state}" | get_json_int last_apply_epoch)"

  local recommend
  recommend="$(should_recommend "${pending}" "${security}" "${last_apply}")"

  local payload
  payload="$(python3 - <<PY
import json,sys,time,os
pending=int("${pending}")
security=int("${security}")
recommend=int("${recommend}")
try:
  state=json.loads(sys.stdin.read() or "{}")
except Exception:
  state={}
state["last_check_iso"]="${iso_now()}"
state["last_check_epoch"]=int(time.time())
state["pending_pkg_count"]=pending
state["security_patch_count"]=security
state["recommend_update"]=bool(recommend)
state["offline_prepared"]=os.path.exists("${OFFLINE_FLAG}")
print(json.dumps(state, indent=2, sort_keys=True))
PY
<<<"${state}")"

  printf '%s\n' "${payload}" | write_state
  printf '%s\n' "${payload}"
}

cmd_download() {
  ensure_root
  load_conf

  log "Refreshing repositories"
  zypper_n refresh

  log "Downloading (staging) updates (zypper dup --download-only)"
  zypper_n dup -y --download-only

  touch "${OFFLINE_FLAG}"

  # Update state
  local state; state="$(read_state)"
  local payload
  payload="$(python3 - <<PY
import json,sys,time,os
try:
  state=json.loads(sys.stdin.read() or "{}")
except Exception:
  state={}
state["last_download_iso"]="${iso_now()}"
state["last_download_epoch"]=int(time.time())
state["offline_prepared"]=True
print(json.dumps(state, indent=2, sort_keys=True))
PY
<<<"${state}")"
  printf '%s\n' "${payload}" | write_state

  log "Updates downloaded. Reboot to apply."
}

cmd_schedule() {
  ensure_root
  touch "${OFFLINE_FLAG}"
  log "Offline apply scheduled (flag set). Reboot to apply."
}

maybe_snapper_pre() {
  # Optional: create rollback snapshot if snapper exists and root config is present.
  if command -v snapper >/dev/null 2>&1; then
    # Find a config named "root" or any default; if none, skip.
    if snapper list-configs >/dev/null 2>&1; then
      if snapper list-configs | awk '{print $1}' | grep -qx "root"; then
        snapper -c root create --type pre --print-number --description "ShiftOS Update Center pre-update" || true
        return 0
      fi
    fi
  fi
  return 0
}

maybe_snapper_post() {
  local pre_id="${1:-}"
  if [[ -n "${pre_id}" ]] && command -v snapper >/dev/null 2>&1; then
    snapper -c root create --type post --pre-number "${pre_id}" --description "ShiftOS Update Center post-update" || true
  fi
}

cmd_apply() {
  ensure_root
  load_conf

  if [[ ! -f "${OFFLINE_FLAG}" ]]; then
    warn "No offline update flag set (${OFFLINE_FLAG}). Running apply anyway."
  fi

  local pre_id=""
  if command -v snapper >/dev/null 2>&1; then
    # capture snapshot id if printed
    pre_id="$(maybe_snapper_pre 2>/dev/null | tail -n1 || true)"
    # if output isn't numeric, ignore
    [[ "${pre_id}" =~ ^[0-9]+$ ]] || pre_id=""
  fi

  log "Applying updates (zypper dup)"
  zypper_n refresh
  zypper_n dup -y

  rm -f "${OFFLINE_FLAG}"

  # Update state
  local state; state="$(read_state)"
  local payload
  payload="$(python3 - <<PY
import json,sys,time
try:
  state=json.loads(sys.stdin.read() or "{}")
except Exception:
  state={}
state["last_apply_iso"]="${iso_now()}"
state["last_apply_epoch"]=int(time.time())
state["offline_prepared"]=False
print(json.dumps(state, indent=2, sort_keys=True))
PY
<<<"${state}")"
  printf '%s\n' "${payload}" | write_state

  maybe_snapper_post "${pre_id}"

  log "Update complete. Reboot may be required."
}

usage() {
  cat <<'EOF'
ShiftOS Update Center (CLI backend)

Usage:
  shift-update-center check        # check + write state.json
  shift-update-center download     # download updates (staging), set offline flag
  shift-update-center schedule     # set offline flag only
  shift-update-center apply        # apply updates (dup), clears offline flag

State:
  /var/lib/shiftos-update-center/state.json
Flag:
  /var/lib/shiftos-update-center/offline_update.flag
Config:
  /etc/shiftos-update-center.conf
EOF
}

main() {
  local cmd="${1:-}"
  case "${cmd}" in
    check) cmd_check ;;
    download) cmd_download ;;
    schedule) cmd_schedule ;;
    apply) cmd_apply ;;
    ""|help|-h|--help) usage ;;
    *) die "Unknown command: ${cmd}" ;;
  esac
}

main "$@"
