#!/bin/bash
set -euo pipefail

TEMP_USER="tempadmin"
KEEP_TEMP_USER=false
USE_TEMP_USER=false
SCHEDULE_MODE=false
STATUS_MODE=false
GID_PROVIDED=false
PATH_RESTRICT="/"
DRY_RUN=false
GRACE_PERIOD=0
LAST_UNIT_FILE="/var/tmp/change_uidgid_last_unit"

print_help() {
cat <<EOF
Usage: $0 --user USERNAME --uid NEW_UID [--gid NEW_GID] [--path DIR]
          [--dry-run] [--use-temp-user] [--keep-temp-user]
          [--schedule] [--grace-period SECONDS] [--status]

Options:
  --user USERNAME         Target username (required)
  --uid NEW_UID           New UID to assign (required)
  --gid NEW_GID           New GID to assign (optional, default: same as UID)
  --path DIR              Restrict file ownership updates (default: '/')
  --dry-run               Show what would be done without changing anything
  --use-temp-user         Create a temp admin user to run changes
  --keep-temp-user        Do not remove the temp admin user afterward
  --schedule              Run change in background via systemd-run
  --grace-period SECONDS  Delay before terminating user sessions (default: 0)
  --status                Show the status of the last scheduled job
  --help                  Show this help message

Examples:
  sudo $0 --user dbi --uid 2222 --schedule --grace-period 10
  sudo $0 --status
EOF
}

# -------- Parse Arguments --------
while [[ $# -gt 0 ]]; do
    case "$1" in
        --user) USERNAME="$2"; shift 2 ;;
        --uid) NEW_UID="$2"; shift 2 ;;
        --gid) NEW_GID="$2"; GID_PROVIDED=true; shift 2 ;;
        --path) PATH_RESTRICT="$2"; shift 2 ;;
        --dry-run) DRY_RUN=true; shift ;;
        --use-temp-user) USE_TEMP_USER=true; shift ;;
        --keep-temp-user) KEEP_TEMP_USER=true; shift ;;
        --schedule) SCHEDULE_MODE=true; shift ;;
        --grace-period) GRACE_PERIOD="$2"; shift 2 ;;
        --status) STATUS_MODE=true; shift ;;
        --help) print_help; exit 0 ;;
        *) echo "❌ Unknown option: $1"; print_help; exit 1 ;;
    esac
done

# -------- Status Mode --------
if $STATUS_MODE; then
    if [[ -f "$LAST_UNIT_FILE" ]]; then
        UNIT_NAME=$(cat "$LAST_UNIT_FILE")
        echo "📊 Status of scheduled job: $UNIT_NAME"
        systemctl status "$UNIT_NAME"
    else
        echo "ℹ️ No scheduled unit found."
    fi
    exit 0
fi

# -------- Validation --------
if [[ -z "${USERNAME:-}" || -z "${NEW_UID:-}" ]]; then
    echo "❌ Missing required parameters."
    print_help
    exit 1
fi

if ! $GID_PROVIDED; then
    NEW_GID="$NEW_UID"
fi

# -------- Schedule Mode --------
if $SCHEDULE_MODE; then
    CMD="$0 --user $USERNAME --uid $NEW_UID --path $PATH_RESTRICT"
    $GID_PROVIDED && CMD+=" --gid $NEW_GID"
    $DRY_RUN && CMD+=" --dry-run"
    [[ "$GRACE_PERIOD" != "0" ]] && CMD+=" --grace-period $GRACE_PERIOD"

    UNIT_NAME="change-uidgid-$USERNAME-$(date +%s)"
    echo "🕑 Scheduling UID/GID change via systemd-run as unit '$UNIT_NAME'..."
    systemd-run --unit="$UNIT_NAME" --description="Change UID/GID of $USERNAME" --no-block $CMD
    echo "$UNIT_NAME" > "$LAST_UNIT_FILE"
    echo "✅ Scheduled. Please log out of '$USERNAME' now. The job will run independently."
    exit 0
fi

# -------- Re-run as Temp User --------
if $USE_TEMP_USER && [[ "${USER:-}" != "$TEMP_USER" ]]; then
    echo "🔧 Creating temporary user '$TEMP_USER'..."
    id "$TEMP_USER" &>/dev/null || useradd -m -s /bin/bash "$TEMP_USER"
    echo "🔁 Re-running script as '$TEMP_USER'..."

    exec sudo -u "$TEMP_USER" "$0" \
        --user "$USERNAME" \
        --uid "$NEW_UID" \
        ${GID_PROVIDED:+--gid "$NEW_GID"} \
        --path "$PATH_RESTRICT" \
        $($DRY_RUN && echo "--dry-run") \
        $($KEEP_TEMP_USER && echo "--keep-temp-user")
fi

# -------- User Exists Check --------
if ! id "$USERNAME" &>/dev/null; then
    echo "❌ User '$USERNAME' does not exist."
    exit 2
fi

OLD_UID=$(id -u "$USERNAME")
OLD_GID=$(id -g "$USERNAME")

echo "🔄 Changing UID:GID for '$USERNAME' from $OLD_UID:$OLD_GID to $NEW_UID:$NEW_GID"
echo "📁 Scanning path: $PATH_RESTRICT"
$DRY_RUN && echo "🧪 Dry-run mode enabled."

# -------- Grace Period --------
if [[ "$GRACE_PERIOD" -gt 0 ]]; then
    echo "⚠️ Will terminate all sessions of '$USERNAME' in $GRACE_PERIOD seconds..."
    sleep "$GRACE_PERIOD"
fi

# -------- Kill Sessions --------
if ! $DRY_RUN; then
    echo "🔪 Terminating user sessions for '$USERNAME'..."
    loginctl terminate-user "$USERNAME" || true
fi

# -------- Change UID/GID --------
if ! $DRY_RUN; then
    groupmod -g "$NEW_GID" "$USERNAME"
    usermod -u "$NEW_UID" -g "$NEW_GID" "$USERNAME"
fi

# -------- Fix Ownership --------
echo "🔍 Updating file ownership..."
if $DRY_RUN; then
    find "$PATH_RESTRICT" -xdev -user "$OLD_UID" -exec echo chown -h "$NEW_UID" {} \;
    find "$PATH_RESTRICT" -xdev -group "$OLD_GID" -exec echo chgrp -h "$NEW_GID" {} \;
else
    find "$PATH_RESTRICT" -xdev -user "$OLD_UID" -exec chown -h "$NEW_UID" {} \;
    find "$PATH_RESTRICT" -xdev -group "$OLD_GID" -exec chgrp -h "$NEW_GID" {} \;
    echo "✅ UID and GID updated successfully for '$USERNAME'."
fi

# -------- Clean Temp User --------
if $USE_TEMP_USER && ! $KEEP_TEMP_USER && [[ "${USER:-}" == "$TEMP_USER" ]]; then
    echo "🧹 Cleaning up temporary user '$TEMP_USER'..."
    userdel -r "$TEMP_USER" || echo "⚠️ Could not delete temp user."
fi