#!/bin/bash
# =============================================================================
# Restore tooling for the borg-backup.sh / dump_db.sh backup system.
# Modes:
#   restore.sh full  [--archive NAME] [--force] [--dry-run]
#   restore.sh db    <db_name> [--archive NAME] [--yes] [--dry-run]
#   restore.sh file  <path-within-target> [--archive NAME] [--dest DIR] [--dry-run]
#   restore.sh --list-archives
# =============================================================================

set -euo pipefail

export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$PATH"
umask 077

# ========================= CONFIGURATION =========================

NAME="borg-2025"
REPO="/home/srv/files/backups/$NAME"
TARGET="/home/srv/files/content"
ARCHIVE_TARGET_PATH="${TARGET#/}"

DB_CONTAINER="mariadb"
DB_START_TIMEOUT=180

BORG_PASSPHRASE_FILE="${BORG_PASSPHRASE_FILE:-/root/.borg-passphrase}"
ROOT_PASSWORD_FILE="${ROOT_PASSWORD_FILE:-/root/.mariadb-root.pw}"
DUMP_SUBDIR="mariadb/dump"

# Same lockfile borg-backup.sh takes (via flock -n 9) before touching $TARGET
# or the repo, so a restore and the nightly backup cron job can never run
# concurrently against each other.
LOCKFILE="${LOCKFILE:-/var/lock/borg-backup.lock}"

LOGDIR="${RESTORE_LOGDIR:-/var/log/borg}"
mkdir -p "$LOGDIR" 2>/dev/null || LOGDIR="/tmp"
LOGFILE="$LOGDIR/restore-$(date +%Y-%m-%d-%H%M%S).log"

export BORG_REPO="$REPO"
export BORG_PASSCOMMAND="cat $BORG_PASSPHRASE_FILE"

# =================================================================

log() {
    local line
    line="[$(date '+%F %T')] $*"
    echo "$line"
    echo "$line" >> "$LOGFILE" 2>/dev/null || true
}

step() { echo; echo "=== $* ==="; }

die() { log "ERROR: $*"; exit 1; }

run_cmd() {
    log "[RUN] $*"
    "$@"
}

acquire_lock() {
    exec 9>"$LOCKFILE"
    if ! flock -n 9; then
        die "another backup or restore is already running (lock held on $LOCKFILE)"
    fi
}

preflight() {
    [[ -r "$BORG_PASSPHRASE_FILE" ]] || die "passphrase file not readable: $BORG_PASSPHRASE_FILE (chmod 600 it or check the path)"
    borg info --lock-wait 60 >/dev/null 2>&1 || die "cannot reach borg repo $REPO (check the passphrase file, permissions, and that the repo exists)"
}

resolve_archive() {
    # Deliberately does NOT filter by hostname (unlike borg-backup.sh's
    # ARCHIVE_GLOB="$(hostname -s)-*" used for prune/list): disaster recovery
    # may need to run from a different host than the one that made the
    # backup, so any archive in the repo is a valid restore candidate.
    if [[ -n "$ARCHIVE_OVERRIDE" ]]; then
        RESOLVED_ARCHIVE="$ARCHIVE_OVERRIDE"
        return 0
    fi
    RESOLVED_ARCHIVE="$(borg list --short --lock-wait 60 2>/dev/null | tail -n1)"
    [[ -n "$RESOLVED_ARCHIVE" ]] || die "no archives found in repo $REPO"
}

list_archives() {
    borg list --lock-wait 60
}

extract_path() {
    local rel_path="$1" dest_dir="$2" archive="$3"
    mkdir -p "$dest_dir"
    ( cd "$dest_dir" && run_cmd borg extract --lock-wait 600 "${REPO}::${archive}" "${ARCHIVE_TARGET_PATH}/${rel_path}" )
    echo "${dest_dir%/}/${ARCHIVE_TARGET_PATH}/${rel_path}"
}

cmd_file() {
    local rel_path="$1" dest
    resolve_archive
    dest="${DEST:-/tmp/restore-file-$$}"
    step "Restoring '$rel_path' from archive $RESOLVED_ARCHIVE into $dest"
    if [[ "$DRY_RUN" == true ]]; then
        log "[DRY-RUN] would extract ${ARCHIVE_TARGET_PATH}/${rel_path} from ${REPO}::${RESOLVED_ARCHIVE} into $dest"
        return 0
    fi
    preflight
    local final
    final="$(extract_path "$rel_path" "$dest" "$RESOLVED_ARCHIVE" | tail -n1)"
    log "Restored file available at: $final"
}

detect_client() {
    if docker exec "$DB_CONTAINER" sh -c 'command -v mariadb' >/dev/null 2>&1; then
        CLIENT_BIN="mariadb"
    else
        CLIENT_BIN="mysql"
    fi
}

get_root_creds() {
    if [[ -n "${MYSQL_ROOT_PASSWORD:-}" ]]; then
        DB_PASS="$MYSQL_ROOT_PASSWORD"
    elif [[ -r "$ROOT_PASSWORD_FILE" ]]; then
        DB_PASS="$(< "$ROOT_PASSWORD_FILE")"
    else
        die "no root DB password: set MYSQL_ROOT_PASSWORD or create $ROOT_PASSWORD_FILE (chmod 600)"
    fi
}

confirm_or_abort() {
    local dbname="$1" typed
    [[ "$YES" == true ]] && return 0
    echo "This will DROP/overwrite database '$dbname'. Type the database name to confirm:"
    read -r typed
    [[ "$typed" == "$dbname" ]] || die "confirmation did not match '$dbname' - aborting"
}

restore_single_db() {
    local sqlfile="$1" dbname="$2"
    [[ -s "$sqlfile" ]] || die "dump file missing or empty: $sqlfile"
    run_cmd docker exec -e MYSQL_PWD="$DB_PASS" "$DB_CONTAINER" \
        "$CLIENT_BIN" -u root -e "CREATE DATABASE IF NOT EXISTS \`$dbname\`;"
    log "[RUN] docker exec -i ... $CLIENT_BIN -u root $dbname < $sqlfile"
    docker exec -i -e MYSQL_PWD="$DB_PASS" "$DB_CONTAINER" \
        "$CLIENT_BIN" -u root "$dbname" < "$sqlfile"
}

cmd_db() {
    local dbname="$1" dumpfile dest
    resolve_archive
    step "Restoring database '$dbname' from archive $RESOLVED_ARCHIVE"
    if [[ "$DRY_RUN" == true ]]; then
        log "[DRY-RUN] would extract ${ARCHIVE_TARGET_PATH}/${DUMP_SUBDIR}/${dbname}.sql from ${REPO}::${RESOLVED_ARCHIVE}"
        log "[DRY-RUN] would DROP/recreate database '$dbname' and import the dump using root credentials"
        return 0
    fi
    confirm_or_abort "$dbname"
    preflight
    dest="/tmp/restore-db-$$"
    dumpfile="$(extract_path "${DUMP_SUBDIR}/${dbname}.sql" "$dest" "$RESOLVED_ARCHIVE" | tail -n1)"
    detect_client
    get_root_creds
    restore_single_db "$dumpfile" "$dbname"
    log "Database '$dbname' restored from $dumpfile"
    rm -rf "$dest"
}

container_running() {
    [[ "$(docker inspect -f '{{.State.Running}}' "$1" 2>/dev/null || echo false)" == "true" ]]
}

container_health() {
    docker inspect -f '{{if .State.Health}}{{.State.Health.Status}}{{else}}none{{end}}' \
        "$1" 2>/dev/null || echo unknown
}

wait_for_container() {
    local name="$1" deadline=$((SECONDS + DB_START_TIMEOUT)) health
    while (( SECONDS < deadline )); do
        if container_running "$name"; then
            health="$(container_health "$name")"
            case "$health" in
                healthy|none) log "$name is up (health: $health)"; return 0 ;;
                unhealthy)    log "WARNING: $name reports unhealthy" ;;
            esac
        fi
        sleep 3
    done
    return 1
}

start_db() {
    docker start "$DB_CONTAINER" >/dev/null 2>&1 || true
    wait_for_container "$DB_CONTAINER"
}

cmd_full() {
    local dumpdir f dbname
    resolve_archive
    step "Full restore from archive $RESOLVED_ARCHIVE into $TARGET"

    if [[ -d "$TARGET" ]] && [[ -n "$(ls -A "$TARGET" 2>/dev/null)" ]] && [[ "$FORCE" != true ]]; then
        die "$TARGET is not empty - pass --force to overwrite (existing data will be replaced)"
    fi

    if [[ "$DRY_RUN" == true ]]; then
        log "[DRY-RUN] would extract full ${ARCHIVE_TARGET_PATH} tree from ${REPO}::${RESOLVED_ARCHIVE} into $TARGET"
        log "[DRY-RUN] would start $DB_CONTAINER and wait for it to become healthy"
        log "[DRY-RUN] would restore every *.sql dump under ${DUMP_SUBDIR}/ using root credentials"
        return 0
    fi

    preflight

    # rm -rf "$TARGET"/* leaves dotfiles behind (stale .nobackup markers, app
    # state) mixed in with the restored tree; find -delete removes everything.
    if [[ -d "$TARGET" ]] && [[ -n "$(ls -A "$TARGET" 2>/dev/null)" ]]; then
        find "${TARGET:?}" -mindepth 1 -delete
    fi
    mkdir -p "$TARGET"

    # Borg records the absolute path it was given at create time, so
    # extracting from / with the leading-slash-stripped path recreates the
    # tree directly at $TARGET - no staging copy, no doubled disk usage.
    ( cd / && run_cmd borg extract --lock-wait 600 "${REPO}::${RESOLVED_ARCHIVE}" "${ARCHIVE_TARGET_PATH}" )

    dumpdir="${TARGET}/${DUMP_SUBDIR}"
    [[ -d "$dumpdir" ]] || die "no dump directory found after extract: $dumpdir"

    # The container must be running before any docker exec against it - on a
    # freshly rebuilt server it's created but stopped, so start it first.
    step "Starting $DB_CONTAINER"
    start_db || die "CRITICAL: $DB_CONTAINER did not come up after extract"

    detect_client
    get_root_creds

    if [[ -f "${dumpdir}/00-users-and-grants.sql" ]]; then
        step "Restoring users and grants"
        run_cmd docker exec -i -e MYSQL_PWD="$DB_PASS" "$DB_CONTAINER" \
            "$CLIENT_BIN" -u root < "${dumpdir}/00-users-and-grants.sql"
    fi

    for f in "$dumpdir"/*.sql; do
        [[ -e "$f" ]] || continue
        dbname="$(basename "$f" .sql)"
        [[ "$dbname" == "00-users-and-grants" ]] && continue
        restore_single_db "$f" "$dbname"
    done

    log "Full restore complete from archive $RESOLVED_ARCHIVE"
}

usage() {
    cat <<'EOF'
Usage:
  restore.sh full  [--archive NAME] [--force] [--dry-run]
  restore.sh db    <db_name> [--archive NAME] [--yes] [--dry-run]
  restore.sh file  <path-within-target> [--archive NAME] [--dest DIR] [--dry-run]
  restore.sh --list-archives
  restore.sh -h | --help
EOF
}

DRY_RUN=false
FORCE=false
YES=false
ARCHIVE_OVERRIDE=""
DEST=""

parse_common_flags() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --archive) ARCHIVE_OVERRIDE="$2"; shift 2 ;;
            --force) FORCE=true; shift ;;
            --yes) YES=true; shift ;;
            --dry-run) DRY_RUN=true; shift ;;
            --dest) DEST="$2"; shift 2 ;;
            *) die "unknown flag: $1" ;;
        esac
    done
}

main() {
    local cmd="${1:-}"
    case "$cmd" in
        "")
            usage
            exit 1
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        --list-archives)
            shift
            parse_common_flags "$@"
            list_archives
            ;;
        full)
            shift
            parse_common_flags "$@"
            acquire_lock
            cmd_full
            ;;
        db)
            shift
            local dbname="${1:-}"
            [[ -n "$dbname" ]] || die "db: missing <db_name>"
            shift
            parse_common_flags "$@"
            acquire_lock
            cmd_db "$dbname"
            ;;
        file)
            shift
            local relpath="${1:-}"
            [[ -n "$relpath" ]] || die "file: missing <path>"
            shift
            parse_common_flags "$@"
            acquire_lock
            cmd_file "$relpath"
            ;;
        *)
            usage
            die "unknown command: $cmd"
            ;;
    esac
}

main "$@"