#!/bin/bash

set -e

TIMESTAMP=$(date +%Y%m%d-%H%M%S)
BACKUP_FILE="system-backup-${TIMESTAMP}.json"

show_help() {
    echo "Usage: $0 [OPTION]"
    echo ""
    echo "Options:"
    echo "  -b, --backup [filename]    Create backup (default: system-backup-TIMESTAMP.json)"
    echo "  -r, --restore <filename>   Restore from backup file"
    echo "  -h, --help                 Show this help message"
}

backup() {
    local filename="${1:-$BACKUP_FILE}"
    
    echo "Creating backup..."
    
    # Get package lists (filtering out any integers/counts)
    local pacman_pkgs=$(paru -Qqen | grep -v '^[0-9]*$')
    local aur_pkgs=$(paru -Qqem | grep -v '^[0-9]*$')
    
    # Get systemd units (filter headers, footers, and integers)
    local system_units=$(systemctl list-unit-files --state=enabled --no-pager --no-legend | \
        awk '{print $1}' | grep -E '\.(service|timer|socket|target|path|mount)$')
    local user_units=$(systemctl --user list-unit-files --state=enabled --no-pager --no-legend | \
        awk '{print $1}' | grep -E '\.(service|timer|socket|target|path|mount)$')
    
    # Get flatpak list
    local flatpak_pkgs=""
    if command -v flatpak &> /dev/null; then
        flatpak_pkgs=$(flatpak list --app --columns=application | tail -n +1 | grep -v '^[0-9]*$')
    fi
    
    # Build JSON
    jq -n \
        --arg timestamp "$TIMESTAMP" \
        --arg pacman "$pacman_pkgs" \
        --arg aur "$aur_pkgs" \
        --arg system "$system_units" \
        --arg user "$user_units" \
        --arg flatpak "$flatpak_pkgs" \
        '{
            timestamp: $timestamp,
            packages: {
                pacman: ($pacman | split("\n") | map(select(length > 0))),
                aur: ($aur | split("\n") | map(select(length > 0)))
            },
            systemd: {
                system: ($system | split("\n") | map(select(length > 0))),
                user: ($user | split("\n") | map(select(length > 0)))
            },
            flatpak: ($flatpak | split("\n") | map(select(length > 0)))
        }' > "$filename"
    
    echo "Backup created: $filename"
}

restore() {
    local filename="$1"
    
    if [[ -z "$filename" ]]; then
        echo "Error: No backup file specified"
        show_help
        exit 1
    fi
    
    if [[ ! -f "$filename" ]]; then
        echo "Error: Backup file $filename not found"
        exit 1
    fi
    
    # Ensure jq is installed
    if ! command -v jq &> /dev/null; then
        echo "Installing jq..."
        sudo pacman -S --noconfirm jq
    fi
    
    echo "Restoring from $filename..."
    
    # Install pacman packages
    echo "Installing official repository packages..."
    jq -r '.packages.pacman[]' "$filename" | paru -S --needed --noconfirm -
    
    # Install AUR packages
    echo "Installing AUR packages..."
    jq -r '.packages.aur[]' "$filename" | paru -S --needed --noconfirm -
    
    # Enable system units
    echo "Enabling system units..."
    jq -r '.systemd.system[]' "$filename" | while read -r unit; do
        [[ -n "$unit" ]] && sudo systemctl enable "$unit" 2>/dev/null || echo "Skipping $unit"
    done
    
    # Enable user units
    echo "Enabling user units..."
    jq -r '.systemd.user[]' "$filename" | while read -r unit; do
        [[ -n "$unit" ]] && systemctl --user enable "$unit" 2>/dev/null || echo "Skipping $unit"
    done
    
    # Install flatpak packages
    if command -v flatpak &> /dev/null; then
        echo "Installing flatpak packages..."
        jq -r '.flatpak[]' "$filename" | while read -r pkg; do
            [[ -n "$pkg" ]] && flatpak install -y "$pkg" 2>/dev/null || echo "Skipping $pkg"
        done
    else
        echo "Flatpak not installed, skipping flatpak packages"
    fi
    
    echo "Restore complete."
}

case "$1" in
    -b|--backup)
        backup "$2"
        ;;
    -r|--restore)
        restore "$2"
        ;;
    -h|--help)
        show_help
        ;;
    *)
        show_help
        exit 1
        ;;
esac