#!/bin/bash

# ==============================================================================
# UNIVERSAL ZFS & SAMBA MANAGER
# ==============================================================================
#
# Creates ZFS datasets + Samba shares wired for cross-platform clients
# (Windows, macOS, Time Machine) via the Apple "fruit" VFS extensions.
#
# Sanitized public version. Companion to:
#   https://www.jonathanbeard.io/blog/2026/06/06/cross-platform-home-storage-nfs-smb.html
#
# EXAMPLES:
#
# 1. Create a standard share:
#    sudo ./manage_server.sh -name=alice-backup -base=tank/data -user=alice -quota=1TB
#
# 2. Abbreviated flags:
#    sudo ./manage_server.sh --n=alice-backup --b=tank/data --u=alice
#
# 3. Create a new user (Linux system + Samba):
#    sudo ./manage_server.sh -create_user -user=bob
#
# 4. Delete a share (removes ZFS dataset & config entry):
#    sudo ./manage_server.sh -delete_share -name=alice-backup -base=tank/data
#
# 5. Erase a user (removes from Linux system + Samba):
#    sudo ./manage_server.sh -erase_user -user=bob
#
# 6. Fix Time Machine global settings:
#    sudo ./manage_server.sh -fix_global
#
# 7. Dry run (test without changes):
#    sudo ./manage_server.sh -name=test -user=alice -base=tank --dry_run
#
# ==============================================================================

# Ensure Root
if [[ $EUID -ne 0 && "$*" != *"--dry_run"* && "$*" != *"--dr"* ]]; then
   echo "Error: This script must be run as root."
   exit 1
fi

# --- DEFAULTS ---
SMB_CONF="/etc/samba/smb.conf"
MNT_BASE="/mnt"
CACHE_MODE="metadata" # Default optimized for backups
DRY_RUN=false
MODE="create_share"   # Default mode
SERVER_IP=$(hostname -I | awk '{print $1}') # Auto-detect IP for instructions

# --- HELPERS ---
usage() {
    echo "Usage: $0 [options]"
    echo ""
    echo "  Share Management:"
    echo "    -name, --n      : Name of the share (e.g., 'alice-backup')"
    echo "    -base, --b      : ZFS parent dataset (e.g., 'tank/data')"
    echo "    -user, --u      : User to assign ownership to"
    echo "    -quota, --q     : Optional size limit (e.g., '500G')"
    echo "    -cache, --c     : ZFS Cache mode (all | metadata | none). Default: metadata"
    echo "    -delete_share   : Deletes the specified Share and ZFS dataset"
    echo ""
    echo "  User Management:"
    echo "    -create_user    : Creates a new System + Samba user"
    echo "    -erase_user     : completely REMOVES a user from System + Samba"
    echo ""
    echo "  System:"
    echo "    -fix_global     : Repairs smb.conf [global] section only"
    echo "    -dry_run, --dr  : Show commands without executing"
    echo ""
    echo "  Example:"
    echo "    $0 -name=backup -user=alice -base=tank/data -quota=1T"
    exit 1
}

# --- ARGUMENT PARSING ---
for arg in "$@"; do
  case "$arg" in
    # Target Configuration
    -name=*|--name=*|--n=*)     SHARE_NAME="${arg#*=}" ;;
    -user=*|--user=*|--u=*)     VALID_USER="${arg#*=}" ;;
    -base=*|--base=*|--b=*)     BASE_DATASET="${arg#*=}" ;;
    -mnt=*|--mnt_base=*|--m=*)  MNT_BASE="${arg#*=}" ;;
    -quota=*|--quota=*|--q=*)   QUOTA_SIZE="${arg#*=}" ;;
    -cache=*|--cache=*|--c=*)   CACHE_MODE="${arg#*=}" ;;

    # Operation Modes
    -fix_global|--fix_global)   MODE="fix_global" ;;
    -delete_share|--delete_share) MODE="delete_share" ;;
    -create_user|--create_user) MODE="create_user" ;;
    -erase_user|--erase_user)   MODE="erase_user" ;;

    # System
    -dry_run|--dry_run|--dr)    DRY_RUN=true ;;
    -help|--help|--h)           usage ;;
  esac
done

# ==============================================================================
# CORE LOGIC
# ==============================================================================

# --- 1. GLOBAL CONFIG REPAIR ---
fix_global_config() {
    local TARGET_FILE=$1
    if ! grep -q "fruit:copyfile = yes" "$TARGET_FILE"; then
        [ "$DRY_RUN" = false ] && echo "[*] Injecting Apple Extensions into [global]..."
        sed -i '/\[global\]/a \   vfs objects = catia fruit streams_xattr\n   fruit:aapl = yes\n   fruit:nfs_aces = no\n   fruit:copyfile = yes\n   fruit:model = MacPro' "$TARGET_FILE"
    fi
}

# --- 2. USER MANAGEMENT ---
manage_user() {
    if [ -z "$VALID_USER" ]; then echo "Error: -user is required."; exit 1; fi

    if [ "$MODE" = "create_user" ]; then
        echo "[*] Creating System User: $VALID_USER"
        if [ "$DRY_RUN" = true ]; then
            echo "DRY RUN: useradd -m -s /bin/bash $VALID_USER"
            echo "DRY RUN: passwd $VALID_USER"
            echo "DRY RUN: smbpasswd -a $VALID_USER"
        else
            if id "$VALID_USER" &>/dev/null; then
                echo "    User already exists. Skipping system creation."
            else
                useradd -m -s /bin/bash "$VALID_USER"
                echo "    Set password for LINUX login:"
                passwd "$VALID_USER"
            fi
            echo "    Set password for SAMBA (Share) access:"
            smbpasswd -a "$VALID_USER"
            smbpasswd -e "$VALID_USER"
        fi

    elif [ "$MODE" = "erase_user" ]; then
        echo "[*] ERASING User: $VALID_USER"
        if [ "$DRY_RUN" = true ]; then
            echo "DRY RUN: smbpasswd -x $VALID_USER"
            echo "DRY RUN: deluser --remove-home $VALID_USER"
        else
            smbpasswd -x "$VALID_USER" 2>/dev/null
            deluser --remove-home "$VALID_USER"
            echo "    User wiped from System and Samba."
        fi
    fi
}

# --- 3. SHARE DELETION ---
delete_share() {
    local CLEAN_NAME="${SHARE_NAME#/}"
    FULL_DATASET_PATH="${BASE_DATASET}/${CLEAN_NAME}"
    echo "[*] Deleting Share: $CLEAN_NAME"

    if [ "$DRY_RUN" = true ]; then
        echo "DRY RUN: zfs destroy -r $FULL_DATASET_PATH"
        echo "DRY RUN: Remove [$CLEAN_NAME] block from $SMB_CONF"
    else
        sed -i "\|^\[$CLEAN_NAME\]|,\|^$|d" "$SMB_CONF"
        zfs destroy -r "$FULL_DATASET_PATH"
        systemctl restart smbd nmbd
        echo "    Share and Dataset deleted."
    fi
}

# --- 4. SHARE CREATION (Main Logic) ---
create_share() {
    # Validations
    if [[ -z "$SHARE_NAME" || -z "$VALID_USER" || -z "$BASE_DATASET" ]]; then
        echo "Error: Missing required fields for creation."
        usage
    fi

    local CLEAN_NAME="${SHARE_NAME#/}"
    FULL_DATASET_PATH="${BASE_DATASET}/${CLEAN_NAME}"
    FINAL_MOUNT_PATH="${MNT_BASE%/}/${CLEAN_NAME}"
    TEMP_CONF=$(mktemp)

    # A. CONFIGURATION RECONSTRUCTION
    cp "$SMB_CONF" "$TEMP_CONF"
    fix_global_config "$TEMP_CONF"

    # Remove old share def using safe delimiter
    sed -i "\|^\[$CLEAN_NAME\]|,\|^$|d" "$TEMP_CONF"

    # Append New Optimized Block
    NEW_BLOCK="\n[$CLEAN_NAME]\n"
    NEW_BLOCK+="   comment = Time Machine / Storage for $VALID_USER\n"
    NEW_BLOCK+="   path = $FINAL_MOUNT_PATH\n"
    NEW_BLOCK+="   valid users = $VALID_USER\n"
    NEW_BLOCK+="   writable = yes\n"
    NEW_BLOCK+="   read only = no\n"
    NEW_BLOCK+="   browsable = yes\n"
    NEW_BLOCK+="   vfs objects = catia fruit streams_xattr acl_xattr\n"
    NEW_BLOCK+="   fruit:time machine = yes\n"
    NEW_BLOCK+="   fruit:metadata = stream\n"
    NEW_BLOCK+="   fruit:posix_rename = yes\n"
    NEW_BLOCK+="   map acl inherit = yes\n"
    NEW_BLOCK+="   store dos attributes = yes"
    echo -e "$NEW_BLOCK" >> "$TEMP_CONF"

    # B. EXECUTION
    if [ "$DRY_RUN" = true ]; then
        echo "--- DRY RUN PREVIEW: ZFS ---"
        echo "zfs create -p -o mountpoint=$FINAL_MOUNT_PATH -o compression=lz4 -o xattr=sa -o acltype=posixacl -o aclmode=passthrough -o primarycache=$CACHE_MODE $FULL_DATASET_PATH"
        echo "--- DRY RUN PREVIEW: SMB.CONF ---"
        tail -n 15 "$TEMP_CONF"
        rm "$TEMP_CONF"
        return
    fi

    # ZFS Operations
    echo "[*] Configuring ZFS Dataset..."
    if ! zfs list "$FULL_DATASET_PATH" >/dev/null 2>&1; then
        zfs create -p \
            -o mountpoint="$FINAL_MOUNT_PATH" \
            -o compression=lz4 \
            -o aclmode=passthrough \
            -o aclinherit=passthrough \
            -o acltype=posixacl \
            -o nbmand=on \
            -o xattr=sa \
            -o primarycache="$CACHE_MODE" \
            -o secondarycache="$CACHE_MODE" \
            ${QUOTA_SIZE:+ -o refquota=$QUOTA_SIZE} \
            "$FULL_DATASET_PATH"
    else
        zfs set acltype=posixacl "$FULL_DATASET_PATH"
        zfs set primarycache="$CACHE_MODE" "$FULL_DATASET_PATH"
        zfs set xattr=sa "$FULL_DATASET_PATH"
        zfs set aclmode=passthrough "$FULL_DATASET_PATH"
    fi

    # Permission Fixes
    mkdir -p "$FINAL_MOUNT_PATH"
    chown "$VALID_USER":"$VALID_USER" "$FINAL_MOUNT_PATH"
    setfacl -d -m u:"$VALID_USER":rwx "$FINAL_MOUNT_PATH"
    setfacl -m u:"$VALID_USER":rwx "$FINAL_MOUNT_PATH"

    # Apply Config
    mv "$TEMP_CONF" "$SMB_CONF"

    # Ensure User is in Samba
    if ! pdbedit -L "$VALID_USER" >/dev/null 2>&1; then
        echo "    User $VALID_USER not in Samba DB. Prompting for password..."
        smbpasswd -a "$VALID_USER"
    fi

    systemctl restart smbd nmbd

    # --- PRETTY PRINT INSTRUCTIONS ---
    echo ""
    echo "=========================================================================="
    echo "                   SHARE CREATED SUCCESSFULLY"
    echo "=========================================================================="
    echo " SHARE NAME:    $CLEAN_NAME"
    echo " USERNAME:      $VALID_USER"
    echo " PASSWORD:      (The Samba password you set earlier)"
    echo " SERVER IP:     $SERVER_IP"
    echo "--------------------------------------------------------------------------"
    echo " [WINDOWS 11 INSTRUCTIONS]"
    echo " 1. Open File Explorer -> Right-click 'This PC' -> 'Map network drive'."
    echo " 2. In the 'Folder' box, type one of the following:"
    echo "      \\\\$(hostname)\\$CLEAN_NAME"
    echo "      \\\\$SERVER_IP\\$CLEAN_NAME"
    echo " 3. Check 'Connect using different credentials' and log in as '$VALID_USER'."
    echo ""
    echo " [macOS / TIME MACHINE INSTRUCTIONS]"
    echo " 1. In Finder, press Cmd+K (Connect to Server)."
    echo " 2. Address: smb://$(hostname).local/$CLEAN_NAME"
    echo " 3. For Time Machine, open Terminal and run:"
    echo "      sudo tmutil setdestination \"smb://$VALID_USER@$(hostname).local/$CLEAN_NAME\""
    echo "=========================================================================="
    echo ""
}

# ==============================================================================
# MAIN SWITCH
# ==============================================================================

case "$MODE" in
    fix_global)
        if [ "$DRY_RUN" = true ]; then echo "Dry run fix global"; exit 0; fi
        fix_global_config "$SMB_CONF"
        systemctl restart smbd nmbd
        echo "Global Config Repaired."
        ;;
    create_user|erase_user)
        manage_user
        ;;
    delete_share)
        delete_share
        ;;
    create_share)
        create_share
        ;;
esac