#!/usr/bin/env bash
# Safe system update with rollback capability

set -euo pipefail

DOTFILES_DIR="$HOME/dotfiles"
SNAPSHOT_DIR="$HOME/.system-snapshots"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
SNAPSHOT_PATH="$SNAPSHOT_DIR/$TIMESTAMP"

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

echo -e "${BLUE}Safe System Update with Rollback Protection${NC}"
echo "=============================================="

# Create snapshot directory
mkdir -p "$SNAPSHOT_PATH"

# Function to create system snapshot
create_snapshot() {
    echo -e "${YELLOW}📸 Creating system snapshot...${NC}"
    
    # Package states
    brew list --versions --formula > "$SNAPSHOT_PATH/brew_formulae.txt"
    brew list --versions --cask > "$SNAPSHOT_PATH/brew_casks.txt"
    mas list > "$SNAPSHOT_PATH/mas_apps.txt" 2>/dev/null || echo "mas not available" > "$SNAPSHOT_PATH/mas_apps.txt"
    
    # Critical configurations
    cp -r "$HOME/.zshrc" "$SNAPSHOT_PATH/" 2>/dev/null || true
    cp -r "$HOME/.gitconfig" "$SNAPSHOT_PATH/" 2>/dev/null || true
    cp -r "$HOME/.ssh/config" "$SNAPSHOT_PATH/" 2>/dev/null || true
    
    # System info
    sw_vers > "$SNAPSHOT_PATH/system_version.txt"
    uname -a > "$SNAPSHOT_PATH/kernel_info.txt"
    
    echo -e "${GREEN}✅ Snapshot created: $SNAPSHOT_PATH${NC}"
}

# Function to test critical tools
test_system_health() {
    echo -e "${YELLOW}🔍 Testing system health...${NC}"
    
    local critical_tools=(git zsh python3 brew)
    local failed_tools=()
    
    for tool in "${critical_tools[@]}"; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            failed_tools+=("$tool")
        fi
    done
    
    if [[ ${#failed_tools[@]} -gt 0 ]]; then
        echo -e "${RED}❌ Critical tools missing: ${failed_tools[*]}${NC}"
        return 1
    fi
    
    # Test shell functionality
    if ! zsh -c "source $HOME/.zshrc && echo 'Shell test passed'" >/dev/null 2>&1; then
        echo -e "${RED}❌ Shell configuration broken${NC}"
        return 1
    fi
    
    echo -e "${GREEN}✅ System health check passed${NC}"
    return 0
}

# Function to rollback
rollback() {
    local snapshot_path="$1"
    echo -e "${YELLOW}🔄 Rolling back to snapshot: $snapshot_path${NC}"
    
    # This is a placeholder - full rollback would need careful implementation
    echo -e "${RED}⚠️  Rollback functionality requires manual implementation${NC}"
    echo "Snapshot available at: $snapshot_path"
    echo "To rollback manually:"
    echo "1. Compare current vs snapshot package lists"
    echo "2. Downgrade specific packages as needed"
    echo "3. Restore configuration files"
}

# Main update process
main() {
    # Pre-update snapshot
    create_snapshot
    
    # Update with error handling
    echo -e "${YELLOW}🔄 Updating packages...${NC}"
    
    if ! brew update; then
        echo -e "${RED}❌ Brew update failed${NC}"
        rollback "$SNAPSHOT_PATH"
        exit 1
    fi
    
    # Capture pre-upgrade state
    brew outdated > "$SNAPSHOT_PATH/outdated_before.txt" || true
    
    if ! brew upgrade; then
        echo -e "${RED}❌ Brew upgrade failed${NC}"
        rollback "$SNAPSHOT_PATH"
        exit 1
    fi
    
    # Test system health post-update
    if ! test_system_health; then
        echo -e "${RED}❌ System health check failed after update${NC}"
        rollback "$SNAPSHOT_PATH"
        exit 1
    fi
    
    # Update lockfile
    "$DOTFILES_DIR/scripts/generate-lockfile.sh"
    
    # Cleanup
    brew cleanup --prune=all
    
    # Success
    echo -e "${GREEN}🎉 Update completed successfully!${NC}"
    echo "Snapshot preserved at: $SNAPSHOT_PATH"
    
    # Clean old snapshots (keep last 10)
    find "$SNAPSHOT_DIR" -maxdepth 1 -type d -name "20*" | sort -r | tail -n +11 | xargs rm -rf 2>/dev/null || true
}

# Run with confirmation
echo -e "${YELLOW}This will update all Homebrew packages with rollback protection.${NC}"
read -p "Continue? (y/N) " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    main
else
    echo "Update cancelled"
fi