#!/bin/bash

# Particle-OS Rollback and Deployment Activation Test Suite
# Tests the atomic rollback and deployment activation functionality

set -euo pipefail

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

# Test configuration
TEST_DIR="/tmp/particle-os-rollback-test"
LOG_FILE="$TEST_DIR/test-rollback-deployment.log"
TEST_PACKAGES=("htop" "curl" "wget" "tree" "vim")
BACKUP_PACKAGES=("nano" "git" "unzip")

# Test counters
TESTS_PASSED=0
TESTS_FAILED=0
TESTS_TOTAL=0

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" | tee -a "$LOG_FILE"
    ((TESTS_PASSED++))
    ((TESTS_TOTAL++))
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
    ((TESTS_FAILED++))
    ((TESTS_TOTAL++))
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$LOG_FILE"
}

# Test helper functions
assert_command_exists() {
    local cmd="$1"
    if command -v "$cmd" >/dev/null 2>&1; then
        log_success "Command exists: $cmd"
        return 0
    else
        log_error "Command not found: $cmd"
        return 1
    fi
}

assert_file_exists() {
    local file="$1"
    if [[ -f "$file" ]]; then
        log_success "File exists: $file"
        return 0
    else
        log_error "File not found: $file"
        return 1
    fi
}

assert_directory_exists() {
    local dir="$1"
    if [[ -d "$dir" ]]; then
        log_success "Directory exists: $dir"
        return 0
    else
        log_error "Directory not found: $dir"
        return 1
    fi
}

assert_command_success() {
    local cmd="$1"
    local description="${2:-Command execution}"
    
    if eval "$cmd" >/dev/null 2>&1; then
        log_success "$description"
        return 0
    else
        log_error "$description failed"
        return 1
    fi
}

# Cleanup function
cleanup() {
    log_info "Cleaning up test environment..."
    
    # Stop any active live overlay
    if sudo apt-layer.sh --live-overlay status 2>/dev/null | grep -q "Active"; then
        log_info "Stopping live overlay..."
        sudo apt-layer.sh --live-overlay stop || true
    fi
    
    # Clean up test packages
    for pkg in "${TEST_PACKAGES[@]}" "${BACKUP_PACKAGES[@]}"; do
        if dpkg -l | grep -q "^ii.*$pkg"; then
            log_info "Removing test package: $pkg"
            sudo apt remove -y "$pkg" || true
        fi
    done
    
    # Clean up test directory
    if [[ -d "$TEST_DIR" ]]; then
        rm -rf "$TEST_DIR"
    fi
    
    log_info "Cleanup completed"
}

# Setup test environment
setup_test_environment() {
    log_info "Setting up test environment..."
    
    # Create test directory
    mkdir -p "$TEST_DIR"
    
    # Check for required tools
    assert_command_exists "apt-layer.sh"
    assert_command_exists "jq"
    assert_command_exists "ostree"
    
    # Initialize apt-layer if needed
    if ! sudo apt-layer.sh status 2>/dev/null | grep -q "Initialized"; then
        log_info "Initializing apt-layer..."
        sudo apt-layer.sh --init
    fi
    
    log_success "Test environment setup completed"
}

# Test 1: Basic OSTree Status and Deployment Database
test_ostree_status() {
    log_info "=== Test 1: OSTree Status and Deployment Database ==="
    
    # Check if deployment database exists
    local deployment_db="/var/lib/particle-os/deployment-db.json"
    assert_file_exists "$deployment_db"
    
    # Check OSTree status command
    assert_command_success "sudo apt-layer.sh ostree status" "OSTree status command"
    
    # Check if status shows current deployment
    local status_output
    status_output=$(sudo apt-layer.sh ostree status 2>/dev/null || echo "")
    if echo "$status_output" | grep -q "Current Deployment"; then
        log_success "OSTree status shows deployment information"
    else
        log_warning "OSTree status may not show deployment information (normal for new systems)"
    fi
    
    log_info "Test 1 completed"
}

# Test 2: Create Initial OSTree Commit
test_create_initial_commit() {
    log_info "=== Test 2: Create Initial OSTree Commit ==="
    
    # Create initial commit with test packages
    log_info "Creating initial OSTree commit with packages: ${TEST_PACKAGES[*]}"
    
    if sudo apt-layer.sh ostree compose install "${TEST_PACKAGES[@]}"; then
        log_success "Initial OSTree commit created successfully"
        
        # Verify commit was created
        local status_output
        status_output=$(sudo apt-layer.sh ostree status 2>/dev/null || echo "")
        if echo "$status_output" | grep -q "Current Deployment"; then
            log_success "Initial commit appears in deployment status"
        else
            log_warning "Initial commit may not appear in status (check manually)"
        fi
    else
        log_error "Failed to create initial OSTree commit"
        return 1
    fi
    
    log_info "Test 2 completed"
}

# Test 3: Create Second OSTree Commit
test_create_second_commit() {
    log_info "=== Test 3: Create Second OSTree Commit ==="
    
    # Create second commit with different packages
    log_info "Creating second OSTree commit with packages: ${BACKUP_PACKAGES[*]}"
    
    if sudo apt-layer.sh ostree compose install "${BACKUP_PACKAGES[@]}"; then
        log_success "Second OSTree commit created successfully"
        
        # Verify we have multiple commits
        local log_output
        log_output=$(sudo apt-layer.sh ostree log 2>/dev/null || echo "")
        local commit_count
        commit_count=$(echo "$log_output" | grep -c "commit" || echo "0")
        
        if [[ "$commit_count" -ge 2 ]]; then
            log_success "Multiple commits found in log ($commit_count commits)"
        else
            log_warning "Expected multiple commits, found: $commit_count"
        fi
    else
        log_error "Failed to create second OSTree commit"
        return 1
    fi
    
    log_info "Test 3 completed"
}

# Test 4: OSTree Log and Diff Functionality
test_ostree_log_and_diff() {
    log_info "=== Test 4: OSTree Log and Diff Functionality ==="
    
    # Test OSTree log command
    assert_command_success "sudo apt-layer.sh ostree log" "OSTree log command"
    
    # Get commit list
    local log_output
    log_output=$(sudo apt-layer.sh ostree log 2>/dev/null || echo "")
    local commits
    commits=($(echo "$log_output" | grep -o "ostree-[0-9]*-[0-9]*" || echo ""))
    
    if [[ ${#commits[@]} -ge 2 ]]; then
        log_success "Found ${#commits[@]} commits for diff testing"
        
        # Test diff between commits
        local commit1="${commits[0]}"
        local commit2="${commits[1]}"
        
        if sudo apt-layer.sh ostree diff "$commit1" "$commit2" >/dev/null 2>&1; then
            log_success "OSTree diff command works between commits"
        else
            log_warning "OSTree diff command may not work as expected"
        fi
    else
        log_warning "Not enough commits for diff testing (found: ${#commits[@]})"
    fi
    
    log_info "Test 4 completed"
}

# Test 5: Rollback Preparation (Without Activation)
test_rollback_preparation() {
    log_info "=== Test 5: Rollback Preparation (Without Activation) ==="
    
    # Get current deployment
    local status_output
    status_output=$(sudo apt-layer.sh ostree status 2>/dev/null || echo "")
    local current_deployment
    current_deployment=$(echo "$status_output" | grep "Current Deployment:" | cut -d: -f2 | tr -d ' ' || echo "")
    
    if [[ -z "$current_deployment" ]]; then
        log_warning "No current deployment found, skipping rollback test"
        return 0
    fi
    
    # Get previous commit for rollback
    local log_output
    log_output=$(sudo apt-layer.sh ostree log 2>/dev/null || echo "")
    local commits
    commits=($(echo "$log_output" | grep -o "ostree-[0-9]*-[0-9]*" || echo ""))
    
    if [[ ${#commits[@]} -ge 2 ]]; then
        local target_commit="${commits[1]}"  # Second commit (previous)
        
        log_info "Preparing rollback to: $target_commit"
        
        if sudo apt-layer.sh ostree rollback "$target_commit"; then
            log_success "Rollback preparation completed successfully"
            
            # Verify pending deployment is set
            local new_status
            new_status=$(sudo apt-layer.sh ostree status 2>/dev/null || echo "")
            if echo "$new_status" | grep -q "Pending Deployment"; then
                log_success "Pending deployment set for rollback"
            else
                log_warning "Pending deployment may not be set correctly"
            fi
        else
            log_error "Rollback preparation failed"
            return 1
        fi
    else
        log_warning "Not enough commits for rollback testing"
    fi
    
    log_info "Test 5 completed"
}

# Test 6: Live Overlay Rollback
test_live_overlay_rollback() {
    log_info "=== Test 6: Live Overlay Rollback ==="
    
    # Start live overlay
    log_info "Starting live overlay..."
    if sudo apt-layer.sh --live-overlay start; then
        log_success "Live overlay started"
        
        # Install a package in live overlay
        log_info "Installing package in live overlay..."
        if sudo apt-layer.sh --live-install "sl"; then
            log_success "Package installed in live overlay"
            
            # Test rollback of live overlay
            log_info "Testing live overlay rollback..."
            if sudo apt-layer.sh --live-overlay rollback; then
                log_success "Live overlay rollback completed"
                
                # Verify package is no longer available
                if ! command -v sl >/dev/null 2>&1; then
                    log_success "Package correctly removed after rollback"
                else
                    log_warning "Package may still be available after rollback"
                fi
            else
                log_error "Live overlay rollback failed"
                return 1
            fi
        else
            log_error "Failed to install package in live overlay"
            return 1
        fi
        
        # Stop live overlay
        sudo apt-layer.sh --live-overlay stop || true
    else
        log_error "Failed to start live overlay"
        return 1
    fi
    
    log_info "Test 6 completed"
}

# Test 7: OSTree Cleanup Functionality
test_ostree_cleanup() {
    log_info "=== Test 7: OSTree Cleanup Functionality ==="
    
    # Test cleanup command
    if sudo apt-layer.sh ostree cleanup; then
        log_success "OSTree cleanup completed successfully"
    else
        log_warning "OSTree cleanup may have failed (check manually)"
    fi
    
    log_info "Test 7 completed"
}

# Test 8: Deployment Activation Simulation
test_deployment_activation_simulation() {
    log_info "=== Test 8: Deployment Activation Simulation ==="
    
    # Check if there's a pending deployment
    local status_output
    status_output=$(sudo apt-layer.sh ostree status 2>/dev/null || echo "")
    local pending_deployment
    pending_deployment=$(echo "$status_output" | grep "Pending Deployment:" | cut -d: -f2 | tr -d ' ' || echo "")
    
    if [[ -n "$pending_deployment" ]]; then
        log_info "Found pending deployment: $pending_deployment"
        log_info "This deployment would activate on next reboot"
        log_success "Deployment activation simulation - pending deployment ready"
        
        # Show what would happen on reboot
        log_info "On reboot, the system would:"
        log_info "1. Activate deployment: $pending_deployment"
        log_info "2. Rollback to previous state"
        log_info "3. Make the rollback the current deployment"
    else
        log_info "No pending deployment found (normal if no rollback was prepared)"
    fi
    
    log_info "Test 8 completed"
}

# Test 9: BootC Alternative Rollback
test_bootc_rollback() {
    log_info "=== Test 9: BootC Alternative Rollback ==="
    
    # Check if bootc-alternative is available
    if ! command -v bootc-alternative.sh >/dev/null 2>&1; then
        log_warning "bootc-alternative.sh not found, skipping BootC rollback test"
        return 0
    fi
    
    # Test BootC status
    if sudo bootc-alternative.sh status >/dev/null 2>&1; then
        log_success "BootC status command works"
        
        # Test BootC rollback (if deployments exist)
        if sudo bootc-alternative.sh rollback >/dev/null 2>&1; then
            log_success "BootC rollback command works"
        else
            log_warning "BootC rollback may not have deployments to rollback"
        fi
    else
        log_warning "BootC status command may not work as expected"
    fi
    
    log_info "Test 9 completed"
}

# Test 10: System Recovery and Safety
test_system_recovery() {
    log_info "=== Test 10: System Recovery and Safety ==="
    
    # Check system integrity
    assert_command_success "sudo apt-layer.sh status" "System status check"
    
    # Check if we can still create new commits
    log_info "Testing system recovery by creating a new commit..."
    if sudo apt-layer.sh ostree compose install "cowsay"; then
        log_success "System recovery test passed - can create new commits"
        
        # Clean up test package
        sudo apt remove -y cowsay || true
    else
        log_error "System recovery test failed - cannot create new commits"
        return 1
    fi
    
    log_info "Test 10 completed"
}

# Main test execution
main() {
    log_info "Starting Particle-OS Rollback and Deployment Activation Test Suite"
    log_info "Test log: $LOG_FILE"
    
    # Setup
    setup_test_environment
    
    # Run tests
    test_ostree_status
    test_create_initial_commit
    test_create_second_commit
    test_ostree_log_and_diff
    test_rollback_preparation
    test_live_overlay_rollback
    test_ostree_cleanup
    test_deployment_activation_simulation
    test_bootc_rollback
    test_system_recovery
    
    # Results
    echo ""
    echo "=== Test Results ==="
    echo "Total Tests: $TESTS_TOTAL"
    echo "Passed: $TESTS_PASSED"
    echo "Failed: $TESTS_FAILED"
    
    if [[ $TESTS_FAILED -eq 0 ]]; then
        log_success "All tests passed! Rollback and deployment activation functionality is working correctly."
        echo ""
        echo "🎉 SUCCESS: Particle-OS rollback and deployment activation is ready for production use!"
        echo ""
        echo "Next steps:"
        echo "1. Test in a real VM environment with actual reboots"
        echo "2. Document any edge cases or warnings found"
        echo "3. Create user guides for rollback procedures"
        exit 0
    else
        log_error "Some tests failed. Please review the log file: $LOG_FILE"
        echo ""
        echo "⚠️  WARNING: Some rollback/deployment tests failed. Review before production use."
        echo ""
        echo "Recommendations:"
        echo "1. Check the log file for specific error details"
        echo "2. Verify system dependencies and permissions"
        echo "3. Test in a controlled environment before production"
        exit 1
    fi
}

# Trap cleanup on exit
trap cleanup EXIT

# Check if running as root
if [[ $EUID -eq 0 ]]; then
    log_error "This script should not be run as root. Please run as a regular user with sudo access."
    exit 1
fi

# Check for sudo access
if ! sudo -n true 2>/dev/null; then
    log_error "This script requires sudo access. Please ensure you can run sudo commands."
    exit 1
fi

# Run main function
main "$@"