#!/bin/bash

# Velero Restore Hook Count Verification Script
# Creates restore with hooks and verifies that Restore Item Operations 
# Attempted/Completed correctly reflect the total number of hooks 
# executed across all pods

set -e

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

# Default values
RESTORE_NAME=""
BACKUP_NAME=""
TARGET_NAMESPACE=""
VELERO_NAMESPACE="openshift-adp"
EXPECTED_PODS=2
EXPECTED_PRE_HOOKS=0  # Restore hooks don't use pre-hooks
EXPECTED_POST_HOOKS=2 # 2 post-hooks per pod (1 successful + 1 failed)
EXPECTED_POST_HOOK_FAILURES=1 # 1 failed post-hook per pod
CREATE_RESTORE=false
WAIT_COMPLETION=true
VERBOSE=false

usage() {
    cat << EOF
Usage: $0 -r <restore-name> -b <backup-name> [options]

Options:
    -r, --restore <n>        Name for the Velero restore (required)
    -b, --backup <n>         Name of the backup to restore from (required)
    -t, --target-namespace <ns> Target namespace for restore (optional)
    -n, --namespace <ns>        Velero namespace (default: openshift-adp)
    -p, --pods <count>          Expected number of pods after restore (default: 2)
    --pre-hooks <count>         Expected number of pre hooks per pod (default: 2)
    --post-hooks <count>        Expected number of post hooks per pod (default: 2)
    -c, --create               Create the restore (if not specified, only verify existing)
    --no-wait                  Don't wait for restore completion
    -v, --verbose              Enable verbose output
    -h, --help                 Show this help message

Example:
    # Create restore and verify
    $0 -r mysql-hooks-restore -b mysql-hooks-backup -c

    # Only verify existing restore
    $0 -r mysql-hooks-restore -b mysql-hooks-backup

    # Restore to different namespace
    $0 -r mysql-hooks-restore -b mysql-hooks-backup -t test-app-restored -c
EOF
}

log() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -r|--restore)
            RESTORE_NAME="$2"
            shift 2
            ;;
        -b|--backup)
            BACKUP_NAME="$2"
            shift 2
            ;;
        -t|--target-namespace)
            TARGET_NAMESPACE="$2"
            shift 2
            ;;
        -n|--namespace)
            VELERO_NAMESPACE="$2"
            shift 2
            ;;
        -p|--pods)
            EXPECTED_PODS="$2"
            shift 2
            ;;
        --pre-hooks)
            EXPECTED_PRE_HOOKS="$2"
            shift 2
            ;;
        --post-hooks)
            EXPECTED_POST_HOOKS="$2"
            shift 2
            ;;
        -c|--create)
            CREATE_RESTORE=true
            shift
            ;;
        --no-wait)
            WAIT_COMPLETION=false
            shift
            ;;
        -v|--verbose)
            VERBOSE=true
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            error "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Validate required parameters
if [[ -z "$RESTORE_NAME" ]]; then
    error "Restore name is required. Use -r or --restore option."
    usage
    exit 1
fi

if [[ -z "$BACKUP_NAME" ]]; then
    error "Backup name is required. Use -b or --backup option."
    usage
    exit 1
fi

# Check if kubectl/oc is available
if command -v oc &> /dev/null; then
    KUBECTL_CMD="oc"
elif command -v kubectl &> /dev/null; then
    KUBECTL_CMD="kubectl"
else
    error "Neither 'oc' nor 'kubectl' command found. Please install one of them."
    exit 1
fi

# Check if velero CLI is available
if ! command -v velero &> /dev/null; then
    error "Velero CLI not found. Please install velero CLI."
    exit 1
fi

log "Using kubectl command: $KUBECTL_CMD"
log "Restore name: $RESTORE_NAME"
log "Backup name: $BACKUP_NAME"
log "Velero namespace: $VELERO_NAMESPACE"

# Check if backup exists
if ! velero backup describe "$BACKUP_NAME" --namespace="$VELERO_NAMESPACE" &> /dev/null; then
    error "Backup '$BACKUP_NAME' not found"
    exit 1
fi

create_restore() {
    log "Creating Velero restore with hooks..."
    
    # Build velero restore command
    RESTORE_CMD="velero restore create $RESTORE_NAME --namespace=$VELERO_NAMESPACE --from-backup=$BACKUP_NAME"
    
    # Add target namespace if specified
    if [[ -n "$TARGET_NAMESPACE" ]]; then
        RESTORE_CMD="$RESTORE_CMD --namespace-mappings=*:$TARGET_NAMESPACE"
        log "Restoring to target namespace: $TARGET_NAMESPACE"
    fi
    
    # Add hooks configuration
    RESTORE_CMD="$RESTORE_CMD --restore-volumes=true"
    
    # Add wait option
    if [[ "$WAIT_COMPLETION" == true ]]; then
        RESTORE_CMD="$RESTORE_CMD --wait"
    fi
    
    # Create restore with hooks using custom resource if complex hooks are needed
    if [[ "$EXPECTED_PRE_HOOKS" -gt 0 || "$EXPECTED_POST_HOOKS" -gt 0 ]]; then
        log "Creating restore with custom hook configuration..."
        
        # Create a custom restore resource with hooks
        cat > /tmp/restore-with-hooks.yaml << EOF
apiVersion: velero.io/v1
kind: Restore
metadata:
  name: $RESTORE_NAME
  namespace: $VELERO_NAMESPACE
spec:
  backupName: $BACKUP_NAME
  hooks:
    resources:
    - name: restore-hook-1
      includedNamespaces:
      - "*"
      postHooks:
      - exec:
          command: ["/bin/sh", "-c", "echo 'Restore hook 1: SUCCESSFUL at \$(date)'"]
          onError: Continue
          timeout: 30s
      - exec:
          command: ["/bin/sh", "-c", "echo 'Restore hook 2: SUCCESSFUL at \$(date)'"]
          onError: Continue
          timeout: 30s
  excludedResources:
  - nodes
  - events
  - events.events.k8s.io
  - backups.velero.io
  - restores.velero.io
  - resticrepositories.velero.io
  - csinodes.storage.k8s.io
  - volumeattachments.storage.k8s.io
  - backuprepositories.velero.io
  itemOperationTimeout: 4h0m0s
EOF
        
        # Apply the restore resource
        if $KUBECTL_CMD apply -f /tmp/restore-with-hooks.yaml; then
            success "Restore '$RESTORE_NAME' created with hooks"
        else
            warn "Failed to create restore with custom hooks, trying velero CLI..."
            # Fallback to velero CLI
            log "Executing: $RESTORE_CMD"
            if eval "$RESTORE_CMD"; then
                success "Restore '$RESTORE_NAME' created successfully"
            else
                error "Failed to create restore '$RESTORE_NAME'"
                return 1
            fi
        fi
        
        # Clean up temp file
        rm -f /tmp/restore-with-hooks.yaml
    else
        # Simple restore without custom hooks
        log "Executing: $RESTORE_CMD"
        if eval "$RESTORE_CMD"; then
            success "Restore '$RESTORE_NAME' created successfully"
        else
            error "Failed to create restore '$RESTORE_NAME'"
            return 1
        fi
    fi
    
    # Wait a bit for the restore to start
    sleep 5
}

verify_restore() {
    log "Verifying restore: $RESTORE_NAME"
    
    # Check if restore exists
    if ! velero restore describe "$RESTORE_NAME" --namespace="$VELERO_NAMESPACE" &> /dev/null; then
        error "Restore '$RESTORE_NAME' not found"
        return 1
    fi
    
    # Wait for restore to complete if requested
    if [[ "$WAIT_COMPLETION" == true ]]; then
        log "Waiting for restore to complete..."
        
        # Wait up to 10 minutes for completion
        for i in {1..60}; do
            PHASE=$(velero restore get --namespace="$VELERO_NAMESPACE" | grep "$RESTORE_NAME" | awk '{print $3}' | head -1)
            if [[ "$PHASE" == "Completed" || "$PHASE" == "Failed" || "$PHASE" == "PartiallyFailed" ]]; then
                break
            fi
            log "Restore phase: $PHASE - waiting... ($i/60)"
            sleep 10
        done
    fi
    
    # Get restore details using both velero CLI and kubectl  
    RESTORE_YAML=$($KUBECTL_CMD get restores.velero.io "$RESTORE_NAME" -n "$VELERO_NAMESPACE" -o yaml 2>/dev/null)
    RESTORE_DESCRIBE=$(velero restore describe "$RESTORE_NAME" --namespace="$VELERO_NAMESPACE" --details 2>/dev/null)
    
    if [[ $VERBOSE == true ]]; then
        log "Restore YAML:"
        echo "$RESTORE_YAML"
        echo ""
        log "Velero describe output:"
        echo "$RESTORE_DESCRIBE"
        echo ""
    fi
    
    # Extract hook information from spec
    PRE_HOOKS_COUNT=$(echo "$RESTORE_YAML" | grep -A 20 "preHooks:" | grep -c "exec:" 2>/dev/null || echo "0")
    POST_HOOKS_COUNT=$(echo "$RESTORE_YAML" | grep -A 20 "postHooks:" | grep -c "exec:" 2>/dev/null || echo "0")
    
    # Extract status information
    PHASE=$(echo "$RESTORE_YAML" | grep "phase:" | awk '{print $2}' | tail -1 || echo "Unknown")
    HOOKS_ATTEMPTED=$(echo "$RESTORE_YAML" | grep "hooksAttempted:" | awk '{print $2}' || echo "0")
    RESTORE_OPS_ATTEMPTED=$(echo "$RESTORE_YAML" | grep "restoreItemOperationsAttempted:" | awk '{print $2}' || echo "0")
    RESTORE_OPS_COMPLETED=$(echo "$RESTORE_YAML" | grep "restoreItemOperationsCompleted:" | awk '{print $2}' || echo "0")
    
    # Also try to get from velero describe output
    if [[ "$RESTORE_OPS_ATTEMPTED" == "0" ]]; then
        RESTORE_OPS_ATTEMPTED=$(echo "$RESTORE_DESCRIBE" | grep -i "restore item operations attempted" | awk '{print $NF}' || echo "0")
        RESTORE_OPS_COMPLETED=$(echo "$RESTORE_DESCRIBE" | grep -i "restore item operations completed" | awk '{print $NF}' || echo "0")
    fi
    
    # Get actual pod count from target namespace
    ACTUAL_PODS=0
    if [[ -n "$TARGET_NAMESPACE" ]]; then
        if $KUBECTL_CMD get namespace "$TARGET_NAMESPACE" &> /dev/null; then
            ACTUAL_PODS=$($KUBECTL_CMD get pods -n "$TARGET_NAMESPACE" --no-headers 2>/dev/null | wc -l || echo "0")
        fi
    else
        # Try to detect from backup source
        BACKUP_NAMESPACES=$(echo "$RESTORE_DESCRIBE" | grep -i "included namespaces" -A 5 | grep -v "included namespaces" | head -5)
        for ns in $BACKUP_NAMESPACES; do
            if $KUBECTL_CMD get namespace "$ns" &> /dev/null; then
                PODS_IN_NS=$($KUBECTL_CMD get pods -n "$ns" --no-headers 2>/dev/null | wc -l || echo "0")
                ACTUAL_PODS=$((ACTUAL_PODS + PODS_IN_NS))
            fi
        done
    fi
    
    if [[ "$ACTUAL_PODS" -eq 0 ]]; then
        warn "Could not determine actual pod count, using expected: $EXPECTED_PODS"
        ACTUAL_PODS=$EXPECTED_PODS
    fi
    
    # Calculate expected values
    EXPECTED_TOTAL_HOOKS=$((EXPECTED_PRE_HOOKS + EXPECTED_POST_HOOKS))
    EXPECTED_TOTAL_OPERATIONS=$((EXPECTED_TOTAL_HOOKS * ACTUAL_PODS))
    
    log "=== Hook Configuration ==="
    log "Pre hooks per pod: $PRE_HOOKS_COUNT"
    log "Post hooks per pod: $POST_HOOKS_COUNT"
    log "Actual pods: $ACTUAL_PODS"
    log "Expected total hooks per pod: $EXPECTED_TOTAL_HOOKS"
    log "Expected total operations: $EXPECTED_TOTAL_OPERATIONS"
    
    # Extract status information
    PHASE=$(echo "$RESTORE_YAML" | grep "phase:" | awk '{print $2}' | tail -1 || echo "Unknown")
    RESTORE_OPS_ATTEMPTED=$(echo "$RESTORE_YAML" | grep "restoreItemOperationsAttempted:" | awk '{print $2}' || echo "0")
    RESTORE_OPS_COMPLETED=$(echo "$RESTORE_YAML" | grep "restoreItemOperationsCompleted:" | awk '{print $2}' || echo "0")
    
    # Also try to get from velero describe output
    if [[ "$RESTORE_OPS_ATTEMPTED" == "0" ]]; then
        RESTORE_OPS_ATTEMPTED=$(echo "$RESTORE_DESCRIBE" | grep -i "restore item operations attempted" | awk '{print $NF}' || echo "0")
        RESTORE_OPS_COMPLETED=$(echo "$RESTORE_DESCRIBE" | grep -i "restore item operations completed" | awk '{print $NF}' || echo "0")
    fi
    
    log ""
    log "=== Restore Status ==="
    log "Phase: $PHASE"
    log "Restore Item Operations Attempted: $RESTORE_OPS_ATTEMPTED"
    log "Restore Item Operations Completed: $RESTORE_OPS_COMPLETED"
    
    log ""
    log "=== Verification Results ==="
    
    ISSUES_FOUND=0
    
    # Verify the main bug: operations should equal (hooks per pod) * (number of pods)
    if [[ "$RESTORE_OPS_ATTEMPTED" -ne "$EXPECTED_TOTAL_OPERATIONS" ]]; then
        error "BUG CONFIRMED: Restore Item Operations Attempted is $RESTORE_OPS_ATTEMPTED, expected $EXPECTED_TOTAL_OPERATIONS"
        error "This indicates the bug where hooks are counted as 1 per pod instead of (hooks per pod) * (number of pods)"
        ((ISSUES_FOUND++))
    else
        success "Restore Item Operations Attempted count is correct: $RESTORE_OPS_ATTEMPTED"
    fi
    
    if [[ "$RESTORE_OPS_COMPLETED" -ne "$EXPECTED_TOTAL_OPERATIONS" ]]; then
        error "BUG CONFIRMED: Restore Item Operations Completed is $RESTORE_OPS_COMPLETED, expected $EXPECTED_TOTAL_OPERATIONS"
        ((ISSUES_FOUND++))
    else
        success "Restore Item Operations Completed count is correct: $RESTORE_OPS_COMPLETED"
    fi
    
    # Summary
    log ""
    log "=== Summary ==="
    if [[ $ISSUES_FOUND -eq 0 ]]; then
        success "All hook operation counts are correct!"
        return 0
    else
        error "Found $ISSUES_FOUND issue(s) with hook operation counting"
        
        log ""
        log "=== Expected Behavior (Velero 1.13+ fix) ==="
        log "- Each exec hook in each pod should be counted separately"
        log "- With $ACTUAL_PODS pods and $EXPECTED_TOTAL_HOOKS hooks per pod:"
        log "  * Restore Item Operations Attempted should be: $EXPECTED_TOTAL_OPERATIONS"
        log "  * Restore Item Operations Completed should be: $EXPECTED_TOTAL_OPERATIONS"
        log ""
        log "=== Current (Buggy) Behavior ==="
        log "- Only 1 operation counted per pod regardless of number of hooks"
        log "- This results in lower counts than expected"
        
        return 1
    fi
}

# Main execution
if [[ "$CREATE_RESTORE" == true ]]; then
    create_restore
    if [[ $? -ne 0 ]]; then
        exit 1
    fi
    
    if [[ "$WAIT_COMPLETION" == true ]]; then
        log "Waiting for restore to complete..."
        sleep 10  # Give some time for status to update
    fi
fi

verify_restore
exit $?
