#!/usr/bin/env bash

# User Quota Setup Script
# This script sets up filesystem quotas for WHP users
# Must be run as root

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    echo -e "${GREEN}✓${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

print_error() {
    echo -e "${RED}✗${NC} $1"
}

print_info() {
    echo -e "${BLUE}ℹ${NC} $1"
}

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    print_error "This script must be run as root"
    exit 1
fi

# Default values
MYSQL_PASSWORD=""
USERS_DIR="/docker/users"
ACTION="setup"

# Function to show usage
show_usage() {
    echo "Usage: $0 [setup|sync|user|remove|status] [options]"
    echo ""
    echo "Actions:"
    echo "  setup                    - Initial quota setup (default)"
    echo "  sync                     - Sync quotas with database"
    echo "  user <username> <mb>     - Set quota for specific user"
    echo "  remove <username>        - Remove quota for user"
    echo "  status                   - Show quota status"
    echo ""
    echo "Options:"
    echo "  -p <password>            - MySQL root password"
    echo "  -d <directory>           - Users directory (default: /docker/users)"
    echo ""
}

# Parse command line arguments
if [ $# -gt 0 ]; then
    ACTION="$1"
    shift
fi

while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--password)
            MYSQL_PASSWORD="$2"
            shift 2
            ;;
        -d|--directory)
            USERS_DIR="$2"
            shift 2
            ;;
        -h|--help)
            show_usage
            exit 0
            ;;
        *)
            break
            ;;
    esac
done

# Get MySQL password from /root/.my.cnf if not provided
get_mysql_password() {
    if [ -z "$MYSQL_PASSWORD" ]; then
        if [ -f "/root/.my.cnf" ]; then
            MYSQL_PASSWORD=$(grep '^password=' /root/.my.cnf | cut -d'=' -f2- | sed 's/^["'\'']//' | sed 's/["'\'']$//')
            if [ -z "$MYSQL_PASSWORD" ]; then
                print_error "Could not read MySQL password from /root/.my.cnf"
                exit 1
            fi
        else
            print_error "MySQL password not provided and /root/.my.cnf not found"
            echo "Use: $0 $ACTION -p <mysql_password>"
            exit 1
        fi
    fi
}

# Test MySQL connection
test_mysql_connection() {
    if ! mysql -h 127.0.0.1 -u root -p"$MYSQL_PASSWORD" -e "SELECT 1;" >/dev/null 2>&1; then
        print_error "Cannot connect to MySQL with provided credentials"
        exit 1
    fi
}

# Detect filesystem for users directory
detect_filesystem() {
    local fs_info=$(df "$USERS_DIR" 2>/dev/null | tail -1)
    local filesystem=$(echo "$fs_info" | awk '{print $1}')
    local mount_point=$(echo "$fs_info" | awk '{print $6}')
    
    echo "$filesystem $mount_point"
}

# Check if quotas are supported
check_quota_support() {
    print_info "Checking quota support..."
    
    # Install quota tools if not present
    if ! command -v quota >/dev/null 2>&1; then
        print_info "Installing quota tools..."
        dnf install -y quota >/dev/null 2>&1 || {
            print_error "Failed to install quota package"
            exit 1
        }
    fi
    
    print_status "Quota tools are available"
    return 0
}

# Enable quotas on filesystem
enable_quotas() {
    print_info "Enabling filesystem quotas..."
    
    local fs_info=($(detect_filesystem))
    local filesystem="${fs_info[0]}"
    local mount_point="${fs_info[1]}"
    
    print_info "Target filesystem: $filesystem mounted at $mount_point"
    
    # Check if quotas are already enabled
    if mount | grep "$mount_point" | grep -qE "(usrquota|uquota)"; then
        print_status "User quotas already enabled on $mount_point"
        return 0
    fi
    
    # Update fstab if needed - handle both device names and UUIDs
    local fstab_line=""
    local fstab_pattern=""
    
    # First try to find by device name
    fstab_line=$(grep "^$filesystem" /etc/fstab)
    if [ -n "$fstab_line" ]; then
        fstab_pattern="^$filesystem"
        print_info "Found device-based fstab entry for $filesystem"
    else
        # Try to find by mount point (for UUID-based entries)
        fstab_line=$(awk -v mount="$mount_point" '$2 == mount {print $0}' /etc/fstab)
        if [ -n "$fstab_line" ]; then
            # Extract the first field (UUID=... or device name)
            local first_field=$(echo "$fstab_line" | awk '{print $1}')
            fstab_pattern="^$(echo "$first_field" | sed 's/[[\.*^$()+?{|]/\\&/g')"
            print_info "Found UUID-based fstab entry for mount point $mount_point"
        fi
    fi
    
    if [ -n "$fstab_line" ]; then
        # Determine the correct quota option based on filesystem type
        local fs_type=$(findmnt -n -o FSTYPE "$mount_point")
        local quota_option="usrquota"
        if [ "$fs_type" = "xfs" ]; then
            quota_option="uquota"
        fi
        
        if ! echo "$fstab_line" | grep -qE "(usrquota|uquota)"; then
            print_info "Adding $quota_option option to /etc/fstab"
            # Create backup
            cp /etc/fstab /etc/fstab.backup.$(date +%Y%m%d-%H%M%S)
            
            # Use awk to safely update the mount options field (4th field) in fstab
            awk -v mount="$mount_point" -v quota_opt="$quota_option" '
            $2 == mount { 
                if ($4 == "defaults") {
                    $4 = "defaults," quota_opt
                } else {
                    $4 = $4 "," quota_opt
                }
            }
            { print }
            ' /etc/fstab > /etc/fstab.tmp && mv /etc/fstab.tmp /etc/fstab
            
            print_info "Reloading systemd configuration..."
            systemctl daemon-reload
            
            # For XFS root filesystem, kernel boot parameters are required
            if [ "$fs_type" = "xfs" ] && [ "$mount_point" = "/" ]; then
                print_info "Configuring kernel boot parameters for XFS root quota support..."
                grubby --args="rootflags=$quota_option" --update-kernel=ALL
                print_info "Regenerating initramfs with quota support..."
                dracut --force
                print_warning "XFS root filesystem quotas require a reboot to become active"
                print_info "After reboot, run: $0 sync"
            else
                print_info "Remounting filesystem to enable quotas..."
                mount -o remount "$mount_point" || {
                    print_error "Failed to remount filesystem"
                    exit 1
                }
            fi
        else
            print_status "Quota options already present in /etc/fstab"
        fi
    else
        print_warning "No fstab entry found for $filesystem mounted at $mount_point - quotas may not persist after reboot"
    fi
    
    # Initialize quotas based on filesystem type
    local fs_type=$(findmnt -n -o FSTYPE "$mount_point")
    print_info "Filesystem type: $fs_type"
    
    if [ "$fs_type" = "xfs" ]; then
        print_info "Initializing XFS quotas..."
        # XFS quotas are enabled at mount time, no need for quotacheck/quotaon
        # Verify quotas are working by checking quota status
        if xfs_quota -x -c "state" "$mount_point" 2>/dev/null | grep -q "User quota state on"; then
            print_status "XFS user quotas are enabled"
        else
            print_warning "XFS quotas may need a reboot to become fully active"
        fi
    else
        print_info "Initializing quota database for $fs_type filesystem..."
        quotacheck -cum "$mount_point" >/dev/null 2>&1 || {
            print_warning "Quota check completed with warnings"
        }
        
        # Turn on quotas
        quotaon "$mount_point" >/dev/null 2>&1 || {
            print_warning "Quotas may already be enabled"
        }
    fi
    
    print_status "Quotas enabled on $mount_point"
}

# Set quota for a user
set_user_quota() {
    local username="$1"
    local disk_limit_mb="$2"
    local quota_limit=$((disk_limit_mb * 1024))            # Quota limit in KB
    local soft_limit=$quota_limit                          # Soft and hard limits are the same
    local hard_limit=$quota_limit                          # Soft and hard limits are the same
    
    local fs_info=($(detect_filesystem))
    local mount_point="${fs_info[1]}"
    local fs_type=$(findmnt -n -o FSTYPE "$mount_point")
    
    # Check if user exists
    if ! id "$username" >/dev/null 2>&1; then
        print_error "User $username does not exist"
        return 1
    fi
    
    # Set the quota based on filesystem type
    print_info "Setting quota for $username: ${disk_limit_mb}MB (both soft and hard limits: ${quota_limit}KB)"
    
    if [ "$fs_type" = "xfs" ]; then
        # Use xfs_quota for XFS filesystems with full path
        if ! command -v xfs_quota >/dev/null 2>&1; then
            if [ -x "/usr/sbin/xfs_quota" ]; then
                XFS_QUOTA="/usr/sbin/xfs_quota"
            else
                print_error "xfs_quota command not found in PATH or /usr/sbin/"
                return 1
            fi
        else
            XFS_QUOTA="xfs_quota"
        fi
        
        $XFS_QUOTA -x -c "limit -u bsoft=${soft_limit}k bhard=${hard_limit}k $username" "$mount_point" || {
            print_error "Failed to set XFS quota for $username"
            return 1
        }
    else
        # Use setquota for other filesystems with full path
        if ! command -v setquota >/dev/null 2>&1; then
            if [ -x "/usr/sbin/setquota" ]; then
                SETQUOTA="/usr/sbin/setquota"
            else
                print_error "setquota command not found in PATH or /usr/sbin/"
                return 1
            fi
        else
            SETQUOTA="setquota"
        fi
        
        $SETQUOTA -u "$username" "$soft_limit" "$hard_limit" 0 0 "$mount_point" || {
            print_error "Failed to set quota for $username"
            return 1
        }
    fi
    
    print_status "Quota set for $username"
    return 0
}

# Remove quota for a user
remove_user_quota() {
    local username="$1"
    local fs_info=($(detect_filesystem))
    local mount_point="${fs_info[1]}"
    local fs_type=$(findmnt -n -o FSTYPE "$mount_point")
    
    print_info "Removing quota for $username"
    
    if [ "$fs_type" = "xfs" ]; then
        # Use xfs_quota for XFS filesystems with full path
        if ! command -v xfs_quota >/dev/null 2>&1; then
            if [ -x "/usr/sbin/xfs_quota" ]; then
                XFS_QUOTA="/usr/sbin/xfs_quota"
            else
                print_error "xfs_quota command not found in PATH or /usr/sbin/"
                return 1
            fi
        else
            XFS_QUOTA="xfs_quota"
        fi
        
        $XFS_QUOTA -x -c "limit -u bsoft=0 bhard=0 $username" "$mount_point" || {
            print_error "Failed to remove XFS quota for $username"
            return 1
        }
    else
        # Use setquota for other filesystems with full path
        if ! command -v setquota >/dev/null 2>&1; then
            if [ -x "/usr/sbin/setquota" ]; then
                SETQUOTA="/usr/sbin/setquota"
            else
                print_error "setquota command not found in PATH or /usr/sbin/"
                return 1
            fi
        else
            SETQUOTA="setquota"
        fi
        
        $SETQUOTA -u "$username" 0 0 0 0 "$mount_point" || {
            print_error "Failed to remove quota for $username"
            return 1
        }
    fi
    
    print_status "Quota removed for $username"
}

# Get user quota usage
get_user_quota() {
    local username="$1"
    local fs_info=($(detect_filesystem))
    local mount_point="${fs_info[1]}"
    
    quota -u "$username" "$mount_point" 2>/dev/null | tail -1
}

# Sync quotas with database
sync_quotas_with_database() {
    print_info "Syncing quotas with database..."
    
    get_mysql_password
    test_mysql_connection
    
    # Get users from database with disk allowances
    local users_data=$(mysql -h 127.0.0.1 -u root -p"$MYSQL_PASSWORD" -sN -e "
        SELECT username, max_disk_space 
        FROM whp.client_allowances 
        WHERE max_disk_space > 0
    " 2>/dev/null)
    
    if [ -z "$users_data" ]; then
        print_warning "No users found in database or client_allowances table missing"
        return 0
    fi
    
    local count=0
    while IFS=$'\t' read -r username disk_limit; do
        if [ -n "$username" ] && [ -n "$disk_limit" ]; then
            if set_user_quota "$username" "$disk_limit"; then
                ((count++))
            fi
        fi
    done <<< "$users_data"
    
    print_status "Synced quotas for $count users"
}

# Show quota status
show_quota_status() {
    local fs_info=($(detect_filesystem))
    local filesystem="${fs_info[0]}"
    local mount_point="${fs_info[1]}"
    local fs_type=$(findmnt -n -o FSTYPE "$mount_point")
    
    echo "=== Quota Status ==="
    echo "Filesystem: $filesystem"
    echo "Mount point: $mount_point"
    echo "Filesystem type: $fs_type"
    echo "Users directory: $USERS_DIR"
    echo ""
    
    # Check if quotas are enabled based on filesystem type
    if [ "$fs_type" = "xfs" ]; then
        if xfs_quota -x -c "state" "$mount_point" 2>/dev/null | grep -q "User quota state on"; then
            print_status "XFS user quotas are ENABLED"
            echo ""
            echo "Current quota usage:"
            xfs_quota -x -c "report -u" "$mount_point" 2>/dev/null || print_warning "Could not retrieve XFS quota report"
        else
            print_warning "XFS user quotas are NOT enabled"
        fi
    else
        if quotaon -p "$mount_point" 2>/dev/null | grep -q "user quota on"; then
            print_status "User quotas are ENABLED"
            echo ""
            echo "Current quota usage:"
            repquota -u "$mount_point" 2>/dev/null || print_warning "Could not retrieve quota report"
        else
            print_warning "User quotas are NOT enabled"
        fi
    fi
}

# Main execution
main() {
    echo "WHP User Quota Management"
    echo "========================="
    echo ""
    
    case "$ACTION" in
        setup)
            print_info "Starting initial quota setup..."
            check_quota_support
            enable_quotas
            sync_quotas_with_database
            print_status "Quota setup completed!"
            ;;
        sync)
            sync_quotas_with_database
            ;;
        user)
            if [ $# -lt 2 ]; then
                print_error "Usage: $0 user <username> <disk_limit_mb>"
                exit 1
            fi
            set_user_quota "$1" "$2"
            ;;
        remove)
            if [ $# -lt 1 ]; then
                print_error "Usage: $0 remove <username>"
                exit 1
            fi
            remove_user_quota "$1"
            ;;
        status)
            show_quota_status
            ;;
        *)
            print_error "Unknown action: $ACTION"
            show_usage
            exit 1
            ;;
    esac
}

# Run main function with remaining arguments
main "$@" 