#!/bin/bash
# ==============================================================================================
# ================================= Rsync Stop =================================================
# ==============================================================================================
#
# PURPOSE
# ─────────────────────────────────────────────────────────────────────────────
# Stops rsync intelligently on both local and remote servers. Auto-detects
# running orchestrators and chooses the safest stop strategy. If an
# orchestrator is running, kills only the rsync subprocess so the orchestrator
# exits cleanly after finishing the current share. Use --full-stop to kill
# everything immediately.
#
# ==============================================================================================
# OPERATIONAL MODEL
# ==============================================================================================
#
# Two Stop Modes
#   Default (smart):
#     Detects if an orchestrator (daily/weekly/critical sync) is running.
#     If orchestrator found → kills rsync subprocess only. Orchestrator sees
#     rsync exit → moves to next share or exits cleanly on its own.
#     If no orchestrator → kills rsync directly (standalone rsync.sh run).
#     Cleans stale lock files after kill.
#     Recovers containers left stopped by interrupted rsync (local only).
#
#   --full-stop (nuclear):
#     Kills orchestrator first → then kills rsync.
#     Orchestrator will NOT continue to next share.
#     Use when everything needs to stop immediately.
#
# Orchestrator Detection
#   detect_rsync_parent() scans all lock files to find which running process
#   has rsync as a descendant. No hardcoded list — works for any orchestrator.
#   Returns "script_name:parent_pid" if found, empty if standalone.
#
# Remote Handling
#   Both local and remote handled in one run via SSH.
#   Remote containers left as-is — docker_watchdog.sh handles remote recovery.
#   If remote unreachable → skips remote cleanly, logs warning.
#
# ==============================================================================================
# OPERATIONAL SAFEGUARDS
# ==============================================================================================
#
# Root Required
#   pkill and docker require root.
#
# Single Instance Lock
#   acquire_lock prevents concurrent stop attempts racing each other.
#
# Timeout Protection
#   DOCKER_TIMEOUT (15s) on all docker calls — hung daemon doesn't block.
#   SSH_TIMEOUT (15s) on all remote SSH calls.
#
# SIGTERM → SIGKILL Sequence
#   Orchestrators receive SIGTERM first, SIGKILL only if still running after 2s.
#
# Container Recovery
#   Restarts local containers left stopped by the killed rsync session.
#   Remote containers deferred to docker_watchdog.sh.
#
# ==============================================================================================
# RUNTIME MODES
# ==============================================================================================
#
# rsync_stop.sh
#     Auto-detect orchestrator. Kill rsync-only or full-stop accordingly.
#
# rsync_stop.sh --full-stop
#     Kill orchestrator first, then kill rsync. Nothing continues after this.
#
# rsync_stop.sh --rsync-only
#     Skip container recovery. Used when called by other scripts that handle
#     recovery themselves.
#
# rsync_stop.sh --dry-run
#     Show what would be killed without killing anything.
#
# rsync_stop.sh --status
#     Show local and remote rsync PIDs, running orchestrators, and lock files.
#
# rsync_stop.sh --full-stop --dry-run
#     Preview full-stop sequence without making any changes.
#
# ==============================================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

source "$SCRIPT_DIR/../load_config.sh"

DOCKER_TIMEOUT=15
SSH_TIMEOUT=15

# ── Parse special flags before parse_args ─────────────────────────────────────────────────────
FULL_STOP=false
RSYNC_ONLY_MODE=false
FILTERED_ARGS=()
for arg in "$@"; do
    case "$arg" in
        --full-stop)   FULL_STOP=true ;;
        --rsync-only)  RSYNC_ONLY_MODE=true ;;
        *)             FILTERED_ARGS+=("$arg") ;;
    esac
done

parse_args "${FILTERED_ARGS[@]}"

# ==============================================================================================
# ━━━ Setup ━━━
# ==============================================================================================
if [[ "$EUID" -ne 0 ]]; then
    error "Must be run as root — pkill and docker require root"
    exit 1
fi


acquire_lock

if ! command -v docker &>/dev/null; then
    error "Docker command not found"
    exit 1
fi

detect_hosts

# Soft IP resolution — rsync_stop continues local-only if remote unreachable
REMOTE_REACHABLE=false
REMOTE_SERVER=$(resolve_tailscale_ip "$REMOTE_SERVER_NAME")
if [[ -z "$REMOTE_SERVER" ]]; then
    warn "$REMOTE_SERVER_NAME — cannot resolve Tailscale IP, remote operations will be skipped"
elif timeout "$SSH_TIMEOUT" ping -c1 -W3 "$REMOTE_SERVER" &>/dev/null; then
    REMOTE_REACHABLE=true
    echo "$REMOTE_SERVER_NAME reachable ✅"
else
    warn "$REMOTE_SERVER_NAME unreachable — remote operations will be skipped"
fi

[[ "$DRY_RUN"   == true ]] && warn "DRY RUN — no changes will be made"
[[ "$FULL_STOP" == true ]] && warn "FULL STOP mode — orchestrator + rsync will be killed"

# ==============================================================================================
# ━━━ Status ━━━
# ==============================================================================================
if [[ "$SHOW_STATUS" == true ]]; then
    echo ""
    echo "━━━━━ $ICON_SUMMARY RSYNC STOP STATUS ━━━━━"
    echo "$ICON_HOST  Identity:  $MY_ID ($LOCAL_SERVER_NAME)"
    echo "$ICON_NET   Remote:    $REMOTE_ID ($REMOTE_SERVER_NAME)"
    echo ""

    LOCAL_PIDS=$(pgrep -x rsync 2>/dev/null | tr '\n' ' ')
    echo "  $ICON_SYNC Local rsync PIDs:  ${LOCAL_PIDS:-none}"

    for lockfile in "$LOCK_DIR"/*.lock; do
        [[ -f "$lockfile" ]] || continue
        content=$(cat "$lockfile" 2>/dev/null)
        pid="${content%%:*}"
        name="${content##*:}"
        [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null && \
            echo "  $ICON_RUNNING Lock: $name (PID $pid)"
    done

    if [[ "$REMOTE_REACHABLE" == true ]]; then
        REMOTE_PIDS=$(timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
            -o ConnectTimeout="$SSH_TIMEOUT" \
            root@"$REMOTE_SERVER" "pgrep -x rsync || true" 2>/dev/null | tr '\n' ' ')
        echo "  $ICON_SYNC Remote rsync PIDs: ${REMOTE_PIDS:-none}"
    else
        echo "  $ICON_WARN Remote: unreachable"
    fi
    echo "━━━━━━━━━━━━━━━━━━━━━━━"
    exit 0
fi

# ==============================================================================================
# ── ORCHESTRATOR DETECTION ────────────────────────────────────────────────────────────────────
# ==============================================================================================
# Scans lock files to find which running process has rsync as a descendant.
# No hardcoded script names — detects any orchestrator automatically.

detect_rsync_parent() {
    local rsync_pids
    rsync_pids=$(pgrep -x rsync 2>/dev/null || true)
    [[ -z "$rsync_pids" ]] && echo "" && return

    for lockfile in "$LOCK_DIR"/*.lock; do
        [[ -f "$lockfile" ]] || continue
        local content pid locked_name
        content=$(cat "$lockfile" 2>/dev/null)
        pid="${content%%:*}"
        locked_name="${content##*:}"
        [[ -z "$pid" ]] && continue
        ! kill -0 "$pid" 2>/dev/null && continue
        [[ "$locked_name" == rsync_* ]] && continue

        local all_descendants
        all_descendants=$(pgrep -P "$pid" 2>/dev/null || true)

        while IFS= read -r rsync_pid; do
            [[ -z "$rsync_pid" ]] && continue
            local ppid
            ppid=$(awk '/^PPid:/{print $2}' /proc/"$rsync_pid"/status 2>/dev/null || echo "")
            if echo "$all_descendants" | grep -qw "$rsync_pid" 2>/dev/null || \
               [[ "$ppid" == "$pid" ]]; then
                echo "${locked_name}:${pid}"
                return
            fi
        done <<< "$rsync_pids"
    done
    echo ""
}

detect_rsync_parent_remote() {
    [[ "$REMOTE_REACHABLE" != true ]] && echo "" && return

    timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
        -o ConnectTimeout="$SSH_TIMEOUT" \
        root@"$REMOTE_SERVER" bash << 'REMOTE_SCRIPT' 2>/dev/null
LOCK_DIR="/tmp/unraid_locks"
rsync_pids=$(pgrep -x rsync 2>/dev/null || true)
[[ -z "$rsync_pids" ]] && exit 0
for lockfile in "$LOCK_DIR"/*.lock; do
    [[ -f "$lockfile" ]] || continue
    content=$(cat "$lockfile" 2>/dev/null)
    pid="${content%%:*}"
    locked_name="${content##*:}"
    [[ -z "$pid" ]] && continue
    ! kill -0 "$pid" 2>/dev/null && continue
    [[ "$locked_name" == rsync_* ]] && continue
    all_descendants=$(pgrep -P "$pid" 2>/dev/null || true)
    while IFS= read -r rsync_pid; do
        [[ -z "$rsync_pid" ]] && continue
        ppid=$(awk '/^PPid:/{print $2}' /proc/"$rsync_pid"/status 2>/dev/null || echo "")
        if echo "$all_descendants" | grep -qw "$rsync_pid" 2>/dev/null || \
           [[ "$ppid" == "$pid" ]]; then
            echo "${locked_name}:${pid}"
            exit 0
        fi
    done <<< "$rsync_pids"
done
REMOTE_SCRIPT
}

LOCAL_ORCH=$(detect_rsync_parent)
REMOTE_ORCH=""
[[ "$REMOTE_REACHABLE" == true ]] && REMOTE_ORCH=$(detect_rsync_parent_remote)

# Determine mode
if [[ "$FULL_STOP" == true ]]; then
    MODE="full-stop"
elif [[ -n "$LOCAL_ORCH" ]] || [[ -n "$REMOTE_ORCH" ]]; then
    MODE="rsync-only"
    [[ -n "$LOCAL_ORCH" ]] && \
        warn "Local orchestrator detected: ${LOCAL_ORCH%%:*} — rsync-only mode"
    [[ -n "$REMOTE_ORCH" ]] && \
        warn "Remote orchestrator detected: ${REMOTE_ORCH%%:*} — rsync-only mode"
    warn "Use --full-stop to also kill the orchestrator"
else
    MODE="rsync-only"
    log "No orchestrator detected — killing rsync directly"
fi

# ==============================================================================================
# ── Kill Orchestrators (full-stop only) ───────────────────────────────────────────────────────
# ==============================================================================================
ORCHESTRATORS_KILLED=()
REMOTE_ORCHESTRATORS_KILLED=()

kill_orchestrator() {
    local script_name="$1" pid="$2"
    local lockfile="$LOCK_DIR/${script_name}.lock"

    if [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would kill $script_name (PID $pid)"
        return 0
    fi

    kill -TERM "$pid" 2>/dev/null
    sleep 2
    kill -0 "$pid" 2>/dev/null && kill -KILL "$pid" 2>/dev/null
    sleep 1

    if ! kill -0 "$pid" 2>/dev/null; then
        warn "$script_name stopped (PID $pid) ✅"
        rm -f "$lockfile"
        return 0
    else
        error "Failed to kill $script_name (PID $pid)"
        return 1
    fi
}

if [[ "$MODE" == "full-stop" ]]; then
    echo ""
    echo "━━━ $ICON_STOP Kill Orchestrators ━━━"

    if [[ -n "$LOCAL_ORCH" ]]; then
        local_name="${LOCAL_ORCH%%:*}"
        local_pid="${LOCAL_ORCH##*:}"
        warn "Killing local: $local_name (PID $local_pid)"
        kill_orchestrator "$local_name" "$local_pid" && \
            ORCHESTRATORS_KILLED+=("$local_name")
    else
        log "No local orchestrator running"
    fi

    if [[ "$REMOTE_REACHABLE" == true ]] && [[ -n "$REMOTE_ORCH" ]]; then
        remote_name="${REMOTE_ORCH%%:*}"
        remote_pid="${REMOTE_ORCH##*:}"
        remote_lock="$LOCK_DIR/${remote_name}.lock"
        warn "Killing remote: $remote_name (PID $remote_pid)"
        if [[ "$DRY_RUN" == false ]]; then
            timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
                -o ConnectTimeout="$SSH_TIMEOUT" \
                root@"$REMOTE_SERVER" \
                "kill -TERM '$remote_pid' 2>/dev/null; sleep 2; \
                 kill -0 '$remote_pid' 2>/dev/null && kill -KILL '$remote_pid' 2>/dev/null; \
                 rm -f '$remote_lock'" 2>/dev/null
            warn "Remote $remote_name stopped ✅"
            REMOTE_ORCHESTRATORS_KILLED+=("$remote_name")
        else
            warn "DRY RUN — would kill remote $remote_name (PID $remote_pid)"
        fi
    elif [[ "$REMOTE_REACHABLE" == true ]]; then
        log "No remote orchestrator running"
    fi

    [[ ${#ORCHESTRATORS_KILLED[@]} -gt 0 || \
       ${#REMOTE_ORCHESTRATORS_KILLED[@]} -gt 0 ]] && sleep 3
fi

# ==============================================================================================
# ━━━ Local Rsync ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_STOP Local Rsync ━━━"

LOCAL_KILLED=false
LOCAL_PIDS=$(pgrep -x rsync 2>/dev/null || true)

if [[ -z "$LOCAL_PIDS" ]]; then
    log "No rsync processes running locally"
else
    warn "Found local rsync PIDs: $(echo "$LOCAL_PIDS" | tr '\n' ' ')"
    if [[ "$DRY_RUN" == true ]]; then
        warn "DRY RUN — would kill local rsync"
    else
        pkill -x rsync 2>/dev/null && LOCAL_KILLED=true || \
            warn "pkill returned non-zero — rsync may have already exited"
        [[ "$LOCAL_KILLED" == true ]] && warn "Local rsync killed ✅"
    fi
fi

# Clean stale rsync lock files
for lockfile in "$LOCK_DIR"/rsync_*.lock; do
    [[ -f "$lockfile" ]] || continue
    content=$(cat "$lockfile" 2>/dev/null)
    pid="${content%%:*}"
    if [[ -n "$pid" ]] && ! kill -0 "$pid" 2>/dev/null; then
        log "Cleaning stale lock: $(basename "$lockfile")"
        [[ "$DRY_RUN" == false ]] && rm -f "$lockfile"
    fi
done

# ==============================================================================================
# ━━━ Remote Rsync ━━━
# ==============================================================================================
echo ""
echo "━━━ $ICON_STOP Remote Rsync — $REMOTE_SERVER_NAME ━━━"

REMOTE_KILLED=false

if [[ "$REMOTE_REACHABLE" == false ]]; then
    warn "Skipping — $REMOTE_SERVER_NAME unreachable"
else
    REMOTE_PIDS=$(timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
        -o ConnectTimeout="$SSH_TIMEOUT" \
        root@"$REMOTE_SERVER" "pgrep -x rsync || true" 2>/dev/null || true)

    if [[ -z "$REMOTE_PIDS" ]]; then
        log "No rsync running on $REMOTE_SERVER_NAME"
    else
        warn "Found remote rsync PIDs: $(echo "$REMOTE_PIDS" | tr '\n' ' ')"
        if [[ "$DRY_RUN" == true ]]; then
            warn "DRY RUN — would kill remote rsync"
        else
            timeout "$SSH_TIMEOUT" ssh -i "$SSH_KEY" \
                -o ConnectTimeout="$SSH_TIMEOUT" \
                root@"$REMOTE_SERVER" "pkill -x rsync || true" 2>/dev/null && \
                REMOTE_KILLED=true || \
                warn "Remote pkill returned non-zero — rsync may have already exited"
            [[ "$REMOTE_KILLED" == true ]] && warn "Remote rsync killed ✅"
        fi
    fi
fi

# ==============================================================================================
# ━━━ Container Recovery ━━━
# ==============================================================================================
# Restart local containers left stopped by interrupted rsync.
# Remote containers left for docker_watchdog.sh to recover.
# Skipped with --rsync-only flag (called by other scripts that handle recovery themselves).
CONTAINERS_RESTARTED=()
CONTAINERS_FAILED=()

if [[ "$RSYNC_ONLY_MODE" == false ]] && \
   { [[ "$LOCAL_KILLED" == true ]] || [[ ${#ORCHESTRATORS_KILLED[@]} -gt 0 ]]; }; then

    echo ""
    echo "━━━ $ICON_START Container Recovery ━━━"
    log "Checking profile containers for recovery..."

    declare -A SEEN
    ALL_CONTAINERS=()

    for profile_containers in "${PROFILE_CRITICAL_CONTAINER_NAMES[@]:-}"; do
        read -r -a container_list <<< "$profile_containers"
        for c in "${container_list[@]:-}"; do
            [[ -z "$c" ]] && continue
            if [[ -z "${SEEN[$c]:-}" ]]; then
                SEEN[$c]=1
                ALL_CONTAINERS+=("$c")
            fi
        done
    done

    if [[ ${#ALL_CONTAINERS[@]} -eq 0 ]]; then
        log "No profile containers defined — skipping recovery"
    else
        for c in "${ALL_CONTAINERS[@]}"; do
            STATUS=$(timeout "$DOCKER_TIMEOUT" docker inspect -f \
                '{{.State.Running}}' "$c" 2>/dev/null || echo "unknown")
            case "$STATUS" in
                true)
                    log "$c — running ✅"
                    ;;
                false)
                    warn "$c — stopped — restarting..."
                    if [[ "$DRY_RUN" == true ]]; then
                        warn "DRY RUN — would restart $c"
                    else
                        if timeout "$DOCKER_TIMEOUT" docker start "$c" >/dev/null 2>&1; then
                            warn "$c restarted ✅"
                            CONTAINERS_RESTARTED+=("$c")
                        else
                            error "Failed to restart $c"
                            CONTAINERS_FAILED+=("$c")
                        fi
                    fi
                    ;;
                *)
                    log "$c not found locally — skipping"
                    ;;
            esac
        done
    fi
fi

# ==============================================================================================
# ━━━ Summary ━━━
# ==============================================================================================
echo ""
echo "━━━━━ $ICON_SUMMARY RSYNC STOP SUMMARY ━━━━━"
echo "$ICON_HOST Identity: $MY_ID ($LOCAL_SERVER_NAME)"
echo "$ICON_GEAR Mode:     $MODE"
echo ""

echo "$ICON_HOST Local ($MY_ID):"
[[ ${#ORCHESTRATORS_KILLED[@]} -gt 0 ]] && \
    warn "  Orchestrators killed: ${ORCHESTRATORS_KILLED[*]}"
if [[ "$LOCAL_KILLED" == true ]]; then
    warn "  Rsync killed ✅"
else
    log "  No rsync was running"
fi

echo "$ICON_NET  Remote ($REMOTE_ID — $REMOTE_SERVER_NAME):"
if [[ "$REMOTE_REACHABLE" == false ]]; then
    warn "  Unreachable — skipped"
else
    [[ ${#REMOTE_ORCHESTRATORS_KILLED[@]} -gt 0 ]] && \
        warn "  Orchestrators killed: ${REMOTE_ORCHESTRATORS_KILLED[*]}"
    if [[ "$REMOTE_KILLED" == true ]]; then
        warn "  Rsync killed ✅"
    else
        log "  No rsync was running"
    fi
fi

[[ ${#CONTAINERS_RESTARTED[@]} -gt 0 ]] && \
    warn "$ICON_CONTAINERS Containers recovered: ${CONTAINERS_RESTARTED[*]}"
[[ ${#CONTAINERS_FAILED[@]} -gt 0 ]] && \
    echo "$ICON_ERROR Containers failed to restart: ${CONTAINERS_FAILED[*]}"

echo ""
if [[ "$DRY_RUN" == true ]]; then
    warn "DRY RUN — no changes made"
else
    echo "$ICON_DONE Status: done ✅"
fi
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Notify if anything was actually killed or failed
if [[ "$DRY_RUN" == false ]]; then
    if [[ ${#CONTAINERS_FAILED[@]} -gt 0 ]]; then
        notify "Rsync stop on $(hostname) ($MY_ID) — containers failed to restart: ${CONTAINERS_FAILED[*]}" \
            "Rsync Stop" "warning"
    elif [[ "$LOCAL_KILLED" == true || "$REMOTE_KILLED" == true || \
            ${#ORCHESTRATORS_KILLED[@]} -gt 0 ]]; then
        notify "Rsync stopped on $(hostname) ($MY_ID) — mode: $MODE${CONTAINERS_RESTARTED:+ — recovered: ${CONTAINERS_RESTARTED[*]}}" \
            "Rsync Stop" "warning"
    fi
fi