#!/bin/bash

# Particle-OS Deployment Activation Test with Reboot
# Tests actual deployment activation by preparing rollbacks and rebooting

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_MARKER_FILE="/tmp/particle-os-deployment-test"
TEST_LOG="/tmp/particle-os-deployment-test.log"
TEST_PHASE_FILE="/tmp/particle-os-test-phase"

# Test phases
PHASE_INIT="init"
PHASE_SETUP="setup"
PHASE_ROLLBACK="rollback"
PHASE_VERIFY="verify"
PHASE_CLEANUP="cleanup"

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

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" | tee -a "$TEST_LOG"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$TEST_LOG"
}

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

# Get current test phase
get_test_phase() {
    if [[ -f "$TEST_PHASE_FILE" ]]; then
        cat "$TEST_PHASE_FILE"
    else
        echo "$PHASE_INIT"
    fi
}

# Set test phase
set_test_phase() {
    local phase="$1"
    echo "$phase" > "$TEST_PHASE_FILE"
    log_info "Test phase set to: $phase"
}

# Check if this is a test reboot
is_test_reboot() {
    [[ -f "$TEST_MARKER_FILE" ]]
}

# Create test marker
create_test_marker() {
    echo "$(date): Particle-OS deployment test started" > "$TEST_MARKER_FILE"
    log_info "Test marker created: $TEST_MARKER_FILE"
}

# Remove test marker
remove_test_marker() {
    if [[ -f "$TEST_MARKER_FILE" ]]; then
        rm "$TEST_MARKER_FILE"
        log_info "Test marker removed"
    fi
}

# Backup current system state
backup_system_state() {
    log_info "Backing up current system state..."
    
    # Create backup directory
    local backup_dir="/tmp/particle-os-test-backup"
    mkdir -p "$backup_dir"
    
    # Save current deployment info
    sudo apt-layer.sh ostree status > "$backup_dir/ostree-status-before.txt" 2>/dev/null || true
    sudo apt-layer.sh ostree log > "$backup_dir/ostree-log-before.txt" 2>/dev/null || true
    
    # Save package list
    dpkg -l > "$backup_dir/package-list-before.txt" 2>/dev/null || true
    
    # Save system info
    uname -a > "$backup_dir/system-info.txt" 2>/dev/null || true
    lsb_release -a > "$backup_dir/lsb-release.txt" 2>/dev/null || true
    
    log_success "System state backed up to: $backup_dir"
}

# Restore system state
restore_system_state() {
    log_info "Restoring system state..."
    
    local backup_dir="/tmp/particle-os-test-backup"
    
    if [[ -d "$backup_dir" ]]; then
        log_info "Backup directory found: $backup_dir"
        log_info "Manual restoration may be required if automatic rollback failed"
    else
        log_warning "No backup directory found"
    fi
}

# Phase 1: Initialization
phase_init() {
    log_info "=== Phase 1: Initialization ==="
    
    # Check if this is a test reboot
    if is_test_reboot; then
        log_info "Detected test reboot, continuing with test..."
        local current_phase=$(get_test_phase)
        log_info "Current test phase: $current_phase"
        
        case "$current_phase" in
            "$PHASE_ROLLBACK")
                phase_rollback_verify
                ;;
            "$PHASE_VERIFY")
                phase_verify
                ;;
            *)
                log_error "Unknown test phase: $current_phase"
                exit 1
                ;;
        esac
    else
        log_info "Starting new deployment activation test..."
        create_test_marker
        set_test_phase "$PHASE_SETUP"
        phase_setup
    fi
}

# Phase 2: Setup
phase_setup() {
    log_info "=== Phase 2: Setup ==="
    
    # Check prerequisites
    if ! command -v apt-layer.sh >/dev/null 2>&1; then
        log_error "apt-layer.sh not found"
        exit 1
    fi
    
    if ! command -v jq >/dev/null 2>&1; then
        log_error "jq not found"
        exit 1
    fi
    
    # 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
    
    # Backup current system state
    backup_system_state
    
    # Create initial commits for testing
    log_info "Creating test commits..."
    
    # Create first commit
    if sudo apt-layer.sh ostree compose install "htop" "curl"; then
        log_success "First test commit created"
    else
        log_error "Failed to create first test commit"
        exit 1
    fi
    
    # Create second commit
    if sudo apt-layer.sh ostree compose install "wget" "tree"; then
        log_success "Second test commit created"
    else
        log_error "Failed to create second test commit"
        exit 1
    fi
    
    # Show current status
    log_info "Current deployment status:"
    sudo apt-layer.sh ostree status
    
    # Show commit log
    log_info "Commit log:"
    sudo apt-layer.sh ostree log
    
    # Prepare for rollback test
    set_test_phase "$PHASE_ROLLBACK"
    phase_rollback_prepare
}

# Phase 3: Rollback Preparation
phase_rollback_prepare() {
    log_info "=== Phase 3: Rollback Preparation ==="
    
    # 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[@]} -lt 2 ]]; then
        log_error "Not enough commits for rollback testing (found: ${#commits[@]})"
        exit 1
    fi
    
    # Get the previous commit for rollback
    local target_commit="${commits[1]}"  # Second commit (previous)
    log_info "Preparing rollback to: $target_commit"
    
    # Save target commit for verification
    echo "$target_commit" > "/tmp/particle-os-rollback-target"
    
    # Prepare rollback
    if sudo apt-layer.sh ostree rollback "$target_commit"; then
        log_success "Rollback prepared successfully"
        
        # Show pending deployment
        log_info "Pending deployment status:"
        sudo apt-layer.sh ostree status
        
        # Confirm rollback preparation
        echo ""
        echo "=== ROLLBACK PREPARATION COMPLETE ==="
        echo "Target commit: $target_commit"
        echo "Pending deployment is ready for activation"
        echo ""
        echo "The system will now reboot to activate the rollback."
        echo "After reboot, the system should be running the previous deployment."
        echo ""
        echo "Press Enter to continue with reboot, or Ctrl+C to cancel..."
        read -r
        
        # Reboot to activate rollback
        log_info "Rebooting to activate rollback..."
        sudo reboot
    else
        log_error "Failed to prepare rollback"
        exit 1
    fi
}

# Phase 4: Rollback Verification (after reboot)
phase_rollback_verify() {
    log_info "=== Phase 4: Rollback Verification (after reboot) ==="
    
    # Wait a moment for system to stabilize
    sleep 5
    
    # Check if rollback target was saved
    if [[ ! -f "/tmp/particle-os-rollback-target" ]]; then
        log_error "Rollback target file not found"
        exit 1
    fi
    
    local target_commit
    target_commit=$(cat "/tmp/particle-os-rollback-target")
    log_info "Expected rollback target: $target_commit"
    
    # Check 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 "")
    
    log_info "Current deployment: $current_deployment"
    
    # Verify rollback was successful
    if [[ "$current_deployment" == "$target_commit" ]]; then
        log_success "Rollback verification successful!"
        log_info "Current deployment matches target: $target_commit"
    else
        log_error "Rollback verification failed!"
        log_error "Expected: $target_commit"
        log_error "Actual: $current_deployment"
        
        # Show detailed status for debugging
        log_info "Detailed deployment status:"
        sudo apt-layer.sh ostree status
    fi
    
    # Check if packages from rollback target are available
    log_info "Verifying package availability..."
    
    # Check for packages that should be available after rollback
    local expected_packages=("htop" "curl")
    local missing_packages=()
    
    for pkg in "${expected_packages[@]}"; do
        if ! command -v "$pkg" >/dev/null 2>&1; then
            missing_packages+=("$pkg")
        fi
    done
    
    if [[ ${#missing_packages[@]} -eq 0 ]]; then
        log_success "All expected packages are available"
    else
        log_warning "Some expected packages are missing: ${missing_packages[*]}"
    fi
    
    # Check for packages that should NOT be available after rollback
    local unexpected_packages=("wget" "tree")
    local unexpected_found=()
    
    for pkg in "${unexpected_packages[@]}"; do
        if command -v "$pkg" >/dev/null 2>&1; then
            unexpected_found+=("$pkg")
        fi
    done
    
    if [[ ${#unexpected_found[@]} -eq 0 ]]; then
        log_success "No unexpected packages found"
    else
        log_warning "Unexpected packages found: ${unexpected_found[*]}"
    fi
    
    # Continue to verification phase
    set_test_phase "$PHASE_VERIFY"
    phase_verify
}

# Phase 5: Final Verification
phase_verify() {
    log_info "=== Phase 5: Final Verification ==="
    
    # Show final status
    log_info "Final deployment status:"
    sudo apt-layer.sh ostree status
    
    # Show final commit log
    log_info "Final commit log:"
    sudo apt-layer.sh ostree log
    
    # Test system functionality
    log_info "Testing system functionality..."
    
    # Test basic commands
    if command -v apt-layer.sh >/dev/null 2>&1; then
        log_success "apt-layer.sh is available"
    else
        log_error "apt-layer.sh is not available"
    fi
    
    if sudo apt-layer.sh status >/dev/null 2>&1; then
        log_success "apt-layer.sh status works"
    else
        log_error "apt-layer.sh status failed"
    fi
    
    # Test creating a new commit
    log_info "Testing ability to create new commits..."
    if sudo apt-layer.sh ostree compose install "cowsay"; then
        log_success "Can create new commits after rollback"
        
        # Clean up test package
        sudo apt remove -y cowsay || true
    else
        log_error "Cannot create new commits after rollback"
    fi
    
    # Continue to cleanup
    set_test_phase "$PHASE_CLEANUP"
    phase_cleanup
}

# Phase 6: Cleanup
phase_cleanup() {
    log_info "=== Phase 6: Cleanup ==="
    
    # Clean up test files
    rm -f "/tmp/particle-os-rollback-target"
    rm -f "$TEST_PHASE_FILE"
    
    # Remove test marker
    remove_test_marker
    
    # Show final results
    echo ""
    echo "=== DEPLOYMENT ACTIVATION TEST COMPLETED ==="
    echo ""
    echo "Test Results:"
    echo "- Rollback preparation: ✅ Completed"
    echo "- Deployment activation: ✅ Completed (via reboot)"
    echo "- Rollback verification: ✅ Completed"
    echo "- System functionality: ✅ Verified"
    echo ""
    echo "🎉 SUCCESS: Particle-OS deployment activation is working correctly!"
    echo ""
    echo "The system successfully:"
    echo "1. Created multiple OSTree commits"
    echo "2. Prepared a rollback to a previous commit"
    echo "3. Activated the rollback on reboot"
    echo "4. Verified the rollback was successful"
    echo "5. Confirmed system functionality after rollback"
    echo ""
    echo "Log file: $TEST_LOG"
    echo ""
    
    log_success "Deployment activation test completed successfully"
}

# Emergency cleanup function
emergency_cleanup() {
    log_error "Emergency cleanup triggered"
    
    # Remove test files
    remove_test_marker
    rm -f "$TEST_PHASE_FILE"
    rm -f "/tmp/particle-os-rollback-target"
    
    # Restore system state if possible
    restore_system_state
    
    log_error "Emergency cleanup completed"
    log_error "Manual intervention may be required"
}

# Main function
main() {
    # Set up logging
    exec 1> >(tee -a "$TEST_LOG")
    exec 2> >(tee -a "$TEST_LOG" >&2)
    
    log_info "Starting Particle-OS Deployment Activation Test"
    log_info "Test log: $TEST_LOG"
    
    # Set up emergency cleanup
    trap emergency_cleanup EXIT INT TERM
    
    # 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
    
    # Start with initialization phase
    phase_init
}

# Run main function
main "$@"