#!/bin/bash

# Folder Snapshot Manager
# Compatible with Ubuntu, Debian, AlmaLinux, Rocky Linux, CentOS
# Uses rsync for efficient and fast snapshot creation

# Configuration
SCRIPT_NAME=$(basename "$0")
CONFIG_DIR="$HOME/.config/snapshot_manager"
DEFAULT_TARGET_DIR=$(pwd)
SNAPSHOTS_BASE_DIR="$HOME/.snapshots/snapshots"
MIN_FREE_SPACE_PERCENT=20

# 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 usage
usage() {
    cat << EOF
Usage: $SCRIPT_NAME [OPTIONS] <COMMAND> [ARGUMENTS]

COMMANDS:
    create <name>           Create a new snapshot with given name
    list                    List all available snapshots (JSON output)
    delete <name>           Delete a snapshot by name
    restore <name>          Replace current folder with snapshot
    info <name>             Show snapshot information
    cleanup                 Remove snapshots older than 30 days

OPTIONS:
    -d, --directory <dir>   Target directory to snapshot (required for first run)
    -h, --help             Show this help message

EXAMPLES:
    $SCRIPT_NAME -d /path/to/folder create backup-v1
    $SCRIPT_NAME list
    $SCRIPT_NAME delete backup-v1
    $SCRIPT_NAME restore backup-v1

Note: The script remembers the target directory after first use.
      Snapshots stored in: ~/.snapshots/
      Special snapshot 'initial-backup' stored in: ~/.initial-backup/
EOF
}

# Function to log messages
log() {
    local level=$1
    shift
    local message="$*"
    
    case $level in
        "ERROR")   echo -e "${RED}[ERROR]${NC} $message" >&2 ;;
        "SUCCESS") echo -e "${GREEN}[SUCCESS]${NC} $message" ;;
        "WARNING") echo -e "${YELLOW}[WARNING]${NC} $message" ;;
        "INFO")    echo -e "${BLUE}[INFO]${NC} $message" ;;
        *)         echo "$message" ;;
    esac
}

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to detect OS
detect_os() {
    if [ -f /etc/os-release ]; then
        source /etc/os-release
        OS_ID="$ID"
        OS_VERSION="$VERSION_ID"
        OS_NAME="$NAME"
    else
        log "ERROR" "Unable to detect operating system"
        return 1
    fi
}

# Function to install rsync automatically
install_rsync() {
    detect_os
    
    log "INFO" "Detected OS: $OS_NAME ($OS_ID $OS_VERSION)"
    
    case "$OS_ID" in
        ubuntu)
            case "$OS_VERSION" in
                22.*|24.*|25.*)
                    log "INFO" "Installing rsync on Ubuntu $OS_VERSION..."
                    if sudo apt update >/dev/null 2>&1 && sudo apt install -y rsync >/dev/null 2>&1; then
                        log "SUCCESS" "rsync installed successfully"
                        return 0
                    else
                        log "ERROR" "Failed to install rsync via apt"
                        return 1
                    fi
                    ;;
                *)
                    log "ERROR" "Unsupported Ubuntu version: $OS_VERSION (supported: 22, 24, 25)"
                    return 1
                    ;;
            esac
            ;;
        debian)
            case "$OS_VERSION" in
                11|12)
                    log "INFO" "Installing rsync on Debian $OS_VERSION..."
                    if sudo apt update >/dev/null 2>&1 && sudo apt install -y rsync >/dev/null 2>&1; then
                        log "SUCCESS" "rsync installed successfully"
                        return 0
                    else
                        log "ERROR" "Failed to install rsync via apt"
                        return 1
                    fi
                    ;;
                *)
                    log "ERROR" "Unsupported Debian version: $OS_VERSION (supported: 11, 12)"
                    return 1
                    ;;
            esac
            ;;
        centos)
            case "$OS_VERSION" in
                9*)
                    log "INFO" "Installing rsync on CentOS Stream $OS_VERSION..."
                    if sudo dnf install -y rsync >/dev/null 2>&1; then
                        log "SUCCESS" "rsync installed successfully"
                        return 0
                    else
                        log "ERROR" "Failed to install rsync via dnf"
                        return 1
                    fi
                    ;;
                *)
                    log "ERROR" "Unsupported CentOS version: $OS_VERSION (supported: 9)"
                    return 1
                    ;;
            esac
            ;;
        almalinux)
            case "$OS_VERSION" in
                9*|10*)
                    log "INFO" "Installing rsync on AlmaLinux $OS_VERSION..."
                    if sudo dnf install -y rsync >/dev/null 2>&1; then
                        log "SUCCESS" "rsync installed successfully"
                        return 0
                    else
                        log "ERROR" "Failed to install rsync via dnf"
                        return 1
                    fi
                    ;;
                *)
                    log "ERROR" "Unsupported AlmaLinux version: $OS_VERSION (supported: 9)"
                    return 1
                    ;;
            esac
            ;;
        rocky)
            case "$OS_VERSION" in
                9*|10*)
                    log "INFO" "Installing rsync on Rocky Linux $OS_VERSION..."
                    if sudo dnf install -y rsync >/dev/null 2>&1; then
                        log "SUCCESS" "rsync installed successfully"
                        return 0
                    else
                        log "ERROR" "Failed to install rsync via dnf"
                        return 1
                    fi
                    ;;
                *)
                    log "ERROR" "Unsupported Rocky Linux version: $OS_VERSION (supported: 9)"
                    return 1
                    ;;
            esac
            ;;
        *)
            log "ERROR" "Unsupported operating system: $OS_ID"
            log "INFO" "Supported systems: Ubuntu (22,24,25), Debian (11,12), CentOS/AlmaLinux/Rocky (9)"
            return 1
            ;;
    esac
}

# Function to check dependencies
check_dependencies() {
    local missing_deps=()
    local rsync_missing=false
    
    for cmd in rsync df du find; do
        if ! command_exists "$cmd"; then
            missing_deps+=("$cmd")
            if [ "$cmd" = "rsync" ]; then
                rsync_missing=true
            fi
        fi
    done
    
    # Auto-install rsync if missing
    if [ "$rsync_missing" = true ]; then
        log "WARNING" "rsync not found. Attempting automatic installation..."
        if install_rsync; then
            # Remove rsync from missing deps since we just installed it
            missing_deps=("${missing_deps[@]/rsync}")
        else
            log "ERROR" "Failed to install rsync automatically"
            exit 1
        fi
    fi
    
    # Check for other missing dependencies
    if [ ${#missing_deps[@]} -ne 0 ]; then
        log "ERROR" "Missing dependencies: ${missing_deps[*]}"
        log "INFO" "These are typically pre-installed system utilities (coreutils, findutils)"
        exit 1
    fi
    
    # log "INFO" "All dependencies satisfied"
}

# Function to load configuration
load_config() {
    local config_file="$CONFIG_DIR/config"
    
    if [ -f "$config_file" ]; then
        source "$config_file"
        if [ -n "$TARGET_DIR" ]; then
            DEFAULT_TARGET_DIR="$TARGET_DIR"
        fi
    fi
}

# Function to save configuration
save_config() {
    mkdir -p "$CONFIG_DIR"
    cat > "$CONFIG_DIR/config" << EOF
TARGET_DIR="$DEFAULT_TARGET_DIR"
EOF
}

# Function to check disk space
check_disk_space() {
    local target_dir="$1"
    local snapshots_dir="$2"
    
    # Check space where snapshots will be stored
    local available_space=$(df "$snapshots_dir" 2>/dev/null | awk 'NR==2 {print $(NF-1)}' | sed 's/%//')
    local free_space=$((100 - available_space))
    
    if [ "$free_space" -lt "$MIN_FREE_SPACE_PERCENT" ]; then
        log "ERROR" "Insufficient disk space. Available: ${free_space}%, Required: ${MIN_FREE_SPACE_PERCENT}%"
        return 1
    fi
    
    log "INFO" "Disk space check passed. Available: ${free_space}%"
    return 0
}

# Function to validate snapshot name
validate_snapshot_name() {
    local name="$1"
    
    if [[ ! "$name" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        log "ERROR" "Invalid snapshot name. Use only alphanumeric characters, hyphens, and underscores."
        return 1
    fi
    
    return 0
}

# Function to create snapshot
create_snapshot() {
    local snapshot_name="$1"
    local target_dir="$2"
    local snapshots_dir="$3"
    
    if [ -z "$snapshot_name" ]; then
        log "ERROR" "Snapshot name is required"
        return 1
    fi
    
    if ! validate_snapshot_name "$snapshot_name"; then
        return 1
    fi
    
    if [ ! -d "$target_dir" ]; then
        log "ERROR" "Target directory does not exist: $target_dir"
        return 1
    fi
    
    # Special handling for initial-backup
    local snapshot_path
    if [ "$snapshot_name" = "initial-backup" ]; then
        snapshot_path="$HOME/.initial-backup"
    else
        snapshot_path="$snapshots_dir/$snapshot_name"
    fi
    
    if [ -d "$snapshot_path" ]; then
        log "ERROR" "Snapshot '$snapshot_name' already exists"
        return 1
    fi
    
    if ! check_disk_space "$target_dir" "$(dirname "$snapshot_path")"; then
        return 1
    fi
    
    log "INFO" "Creating snapshot '$snapshot_name' from '$target_dir'..."
    
    # Create snapshot directory
    mkdir -p "$snapshot_path"
    
    # Create metadata
    local created_timestamp=$(date +%s)
    cat > "$snapshot_path/.snapshot_info" << EOF
name=$snapshot_name
source=$target_dir
created=$(date '+%Y-%m-%d %H:%M:%S')
created_timestamp=$created_timestamp
state=in-progress
EOF
    
    # Use rsync for efficient copying with optimizations
    local start_time=$(date +%s)
    if rsync -a --delete --info=progress2 "$target_dir/" "$snapshot_path/data/" 2>/dev/null; then
        local end_time=$(date +%s)
        local duration=$((end_time - start_time))
        
        # Calculate size efficiently
        local size=$(du -sh "$snapshot_path/data" 2>/dev/null | cut -f1)
        
        # Update metadata with completion
        cat > "$snapshot_path/.snapshot_info" << EOF
name=$snapshot_name
source=$target_dir
created=$(date '+%Y-%m-%d %H:%M:%S')
created_timestamp=$created_timestamp
size=$size
duration=${duration}s
state=completed
EOF
        
        log "SUCCESS" "Snapshot '$snapshot_name' created successfully"
        log "INFO" "Size: $size, Duration: ${duration}s"
        return 0
    else
        log "ERROR" "Failed to create snapshot"
        # Update state to failed before cleanup
        sed -i 's/state=in-progress/state=failed/' "$snapshot_path/.snapshot_info" 2>/dev/null
        return 1
    fi
}

# Function to list snapshots
list_snapshots() {
    local snapshots_dir="$1"
    local snapshots=()
    
    # Collect snapshots from regular directory
    if [ -d "$snapshots_dir" ]; then
        while IFS= read -r -d '' snapshot_path; do
            snapshots+=("$snapshot_path")
        done < <(find "$snapshots_dir" -maxdepth 1 -type d -name "*" ! -path "$snapshots_dir" -print0 2>/dev/null)
    fi
    
    # Add initial-backup if it exists
    if [ -d "$HOME/.initial-backup" ]; then
        snapshots+=("$HOME/.initial-backup")
    fi
    
    # Output JSON
    echo -n "["
    local first=true
    
    for snapshot_path in "${snapshots[@]}"; do
        local info_file="$snapshot_path/.snapshot_info"
        
        if [ -f "$info_file" ]; then
            local name=$(grep "^name=" "$info_file" 2>/dev/null | cut -d'=' -f2)
            local timestamp=$(grep "^created_timestamp=" "$info_file" 2>/dev/null | cut -d'=' -f2)
            local state=$(grep "^state=" "$info_file" 2>/dev/null | cut -d'=' -f2)
            
            # Default values if not found
            [ -z "$name" ] && name=$(basename "$snapshot_path")
            [ -z "$timestamp" ] && timestamp=0
            [ -z "$state" ] && state="completed"
            
            if [ "$first" = true ]; then
                first=false
            else
                echo -n ","
            fi
            
            echo -n "{\"name\":\"$name\",\"time\":$timestamp,\"state\":\"$state\"}"
        fi
    done
    
    echo "]"
}

# Function to delete snapshot
delete_snapshot() {
    local snapshot_name="$1"
    local snapshots_dir="$2"
    
    if [ -z "$snapshot_name" ]; then
        log "ERROR" "Snapshot name is required"
        return 1
    fi
    
    # Determine snapshot path
    local snapshot_path
    if [ "$snapshot_name" = "initial-backup" ]; then
        snapshot_path="$HOME/.initial-backup"
    else
        snapshot_path="$snapshots_dir/$snapshot_name"
    fi
    
    if [ ! -d "$snapshot_path" ]; then
        log "ERROR" "Snapshot '$snapshot_name' does not exist"
        return 1
    fi
    
    log "INFO" "Deleting snapshot '$snapshot_name'..."
    
    if rm -rf "$snapshot_path"; then
        log "SUCCESS" "Snapshot '$snapshot_name' deleted successfully"
        return 0
    else
        log "ERROR" "Failed to delete snapshot '$snapshot_name'"
        return 1
    fi
}

# Function to restore snapshot
restore_snapshot() {
    local snapshot_name="$1"
    local target_dir="$2"
    local snapshots_dir="$3"
    
    if [ -z "$snapshot_name" ]; then
        log "ERROR" "Snapshot name is required"
        return 1
    fi
    
    # Determine snapshot path
    local snapshot_path
    if [ "$snapshot_name" = "initial-backup" ]; then
        snapshot_path="$HOME/.initial-backup"
    else
        snapshot_path="$snapshots_dir/$snapshot_name"
    fi
    
    local snapshot_data_path="$snapshot_path/data"
    
    if [ ! -d "$snapshot_path" ]; then
        log "ERROR" "Snapshot '$snapshot_name' does not exist"
        return 1
    fi
    
    if [ ! -d "$snapshot_data_path" ]; then
        log "ERROR" "Snapshot data directory not found: $snapshot_data_path"
        return 1
    fi
    
    log "INFO" "Restoring snapshot '$snapshot_name' to '$target_dir'..."
    
    # Create backup of current state
    local backup_name="auto-backup-$(date +%s)"
    local backup_path="$snapshots_dir/$backup_name"
    
    if [ -d "$target_dir" ]; then
        log "INFO" "Creating automatic backup '$backup_name' of current state..."
        mkdir -p "$backup_path"
        local backup_timestamp=$(date +%s)
        cat > "$backup_path/.snapshot_info" << EOF
name=$backup_name
source=$target_dir
created=$(date '+%Y-%m-%d %H:%M:%S')
created_timestamp=$backup_timestamp
type=auto-backup
state=completed
EOF
        rsync -a --delete "$target_dir/" "$backup_path/data/" >/dev/null 2>&1
    fi
    
    # Perform restoration
    if rsync -a --delete "$snapshot_data_path/" "$target_dir/" 2>/dev/null; then
        log "SUCCESS" "Snapshot '$snapshot_name' restored successfully"
        log "INFO" "Automatic backup created: '$backup_name'"
        return 0
    else
        log "ERROR" "Failed to restore snapshot"
        return 1
    fi
}

# Function to show snapshot info
show_snapshot_info() {
    local snapshot_name="$1"
    local snapshots_dir="$2"
    
    if [ -z "$snapshot_name" ]; then
        log "ERROR" "Snapshot name is required"
        return 1
    fi
    
    # Determine snapshot path
    local snapshot_path
    if [ "$snapshot_name" = "initial-backup" ]; then
        snapshot_path="$HOME/.initial-backup"
    else
        snapshot_path="$snapshots_dir/$snapshot_name"
    fi
    
    local info_file="$snapshot_path/.snapshot_info"
    
    if [ ! -d "$snapshot_path" ]; then
        log "ERROR" "Snapshot '$snapshot_name' does not exist"
        return 1
    fi
    
    if [ -f "$info_file" ]; then
        echo "Snapshot Information:"
        echo "===================="
        cat "$info_file"
        echo ""
        echo "Files: $(find "$snapshot_path/data" -type f 2>/dev/null | wc -l)"
        echo "Directories: $(find "$snapshot_path/data" -type d 2>/dev/null | wc -l)"
    else
        log "WARNING" "No metadata found for snapshot '$snapshot_name'"
    fi
}

# Function to cleanup old snapshots
cleanup_snapshots() {
    local snapshots_dir="$1"
    local days_old=30
    
    if [ ! -d "$snapshots_dir" ]; then
        log "INFO" "No snapshots directory found"
        return 0
    fi
    
    log "INFO" "Cleaning up snapshots older than $days_old days..."
    
    local cleaned=0
    while IFS= read -r -d '' snapshot_path; do
        local snapshot_name=$(basename "$snapshot_path")
        local info_file="$snapshot_path/.snapshot_info"
        
        if [ -f "$info_file" ]; then
            local created_timestamp=$(grep "^created_timestamp=" "$info_file" 2>/dev/null | cut -d'=' -f2)
            local current_timestamp=$(date +%s)
            local age_days=$(( (current_timestamp - created_timestamp) / 86400 ))
            
            if [ "$age_days" -gt "$days_old" ]; then
                log "INFO" "Removing old snapshot '$snapshot_name' (${age_days} days old)"
                rm -rf "$snapshot_path"
                ((cleaned++))
            fi
        fi
    done < <(find "$snapshots_dir" -maxdepth 1 -type d -name "*" ! -path "$snapshots_dir" -print0 2>/dev/null)
    
    log "SUCCESS" "Cleanup completed. Removed $cleaned snapshots."
}

# Main function
main() {
    local target_dir=""
    local command=""
    local snapshot_name=""
    
    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -d|--directory)
                target_dir="$2"
                shift 2
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            create|list|delete|restore|info|cleanup)
                command="$1"
                if [[ "$command" != "list" && "$command" != "cleanup" ]]; then
                    snapshot_name="$2"
                    shift 2
                else
                    shift
                fi
                break
                ;;
            *)
                log "ERROR" "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done
    
    # Check dependencies
    check_dependencies
    
    # Load configuration
    load_config
    
    # Set target directory
    if [ -z "$target_dir" ]; then
        target_dir="$DEFAULT_TARGET_DIR"
    fi
    
    # Validate required parameters
    if [[ "$command" != "list" && "$command" != "cleanup" && -z "$target_dir" ]]; then
        log "ERROR" "Target directory is required. Use -d option or run setup first."
        usage
        exit 1
    fi
    
    if [ -z "$command" ]; then
        log "ERROR" "Command is required"
        usage
        exit 1
    fi
    
    # Save configuration if target directory is specified
    if [ -n "$target_dir" ]; then
        DEFAULT_TARGET_DIR="$target_dir"
        save_config
    fi
    
    # Create snapshots directory if needed
    mkdir -p "$SNAPSHOTS_BASE_DIR"
    
    # Execute command
    case $command in
        create)
            create_snapshot "$snapshot_name" "$target_dir" "$SNAPSHOTS_BASE_DIR"
            ;;
        list)
            list_snapshots "$SNAPSHOTS_BASE_DIR"
            ;;
        delete)
            delete_snapshot "$snapshot_name" "$SNAPSHOTS_BASE_DIR"
            ;;
        restore)
            restore_snapshot "$snapshot_name" "$target_dir" "$SNAPSHOTS_BASE_DIR"
            ;;
        info)
            show_snapshot_info "$snapshot_name" "$SNAPSHOTS_BASE_DIR"
            ;;
        cleanup)
            cleanup_snapshots "$SNAPSHOTS_BASE_DIR"
            ;;
        *)
            log "ERROR" "Invalid command: $command"
            usage
            exit 1
            ;;
    esac
}

# Run main function
main "$@"