#!/bin/bash

# Automated Image Digest Management Script
# Optimized version of generate_image_digest_lock.sh with automation features

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
STACKS_DIR="$PROJECT_ROOT/stacks"
LOCK_FILE="$PROJECT_ROOT/configs/image-digest-lock.yaml"
LOG_FILE="$PROJECT_ROOT/logs/image-update-$(date +%Y%m%d-%H%M%S).log"

# Create directories if they don't exist
mkdir -p "$(dirname "$LOCK_FILE")" "$PROJECT_ROOT/logs"

# Logging function
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$LOG_FILE"
}

# Function to extract images from stack files
extract_images() {
    local stack_file="$1"
    
    # Use yq to extract image names from Docker Compose files
    if command -v yq >/dev/null 2>&1; then
        yq eval '.services[].image' "$stack_file" 2>/dev/null | grep -v "null" || true
    else
        # Fallback to grep if yq is not available
        grep -E "^\s*image:\s*" "$stack_file" | sed 's/.*image:\s*//' | sed 's/\s*$//' || true
    fi
}

# Function to get image digest from registry
get_image_digest() {
    local image="$1"
    local digest=""
    
    # Handle images without explicit tag (assume :latest)
    if [[ "$image" != *":"* ]]; then
        image="${image}:latest"
    fi
    
    log "Fetching digest for $image"
    
    # Try to get digest from Docker registry
    if command -v skopeo >/dev/null 2>&1; then
        digest=$(skopeo inspect "docker://$image" 2>/dev/null | jq -r '.Digest' || echo "")
    else
        # Fallback to docker manifest inspect (requires Docker CLI)
        digest=$(docker manifest inspect "$image" 2>/dev/null | jq -r '.config.digest' || echo "")
    fi
    
    if [[ -n "$digest" && "$digest" != "null" ]]; then
        echo "$digest"
    else
        log "Warning: Could not fetch digest for $image"
        echo ""
    fi
}

# Function to process all stack files and generate lock file
generate_digest_lock() {
    log "Starting automated image digest lock generation"
    
    # Initialize lock file
    cat > "$LOCK_FILE" << 'EOF'
# Automated Image Digest Lock File
# Generated by automated-image-update.sh
# DO NOT EDIT MANUALLY - This file is automatically updated

version: "1.0"
generated_at: "$(date -Iseconds)"
images:
EOF
    
    # Find all stack YAML files
    local stack_files
    stack_files=$(find "$STACKS_DIR" -name "*.yml" -o -name "*.yaml" 2>/dev/null || true)
    
    if [[ -z "$stack_files" ]]; then
        log "No stack files found in $STACKS_DIR"
        return 1
    fi
    
    declare -A processed_images
    local total_images=0
    local successful_digests=0
    
    # Process each stack file
    while IFS= read -r stack_file; do
        log "Processing stack file: $stack_file"
        
        local images
        images=$(extract_images "$stack_file")
        
        if [[ -n "$images" ]]; then
            while IFS= read -r image; do
                [[ -z "$image" ]] && continue
                
                # Skip if already processed
                if [[ -n "${processed_images[$image]:-}" ]]; then
                    continue
                fi
                
                ((total_images++))
                processed_images["$image"]=1
                
                local digest
                digest=$(get_image_digest "$image")
                
                if [[ -n "$digest" ]]; then
                    # Add to lock file
                    cat >> "$LOCK_FILE" << EOF
  "$image":
    digest: "$digest"
    pinned_reference: "${image%:*}@$digest"
    last_updated: "$(date -Iseconds)"
    source_stack: "$(basename "$stack_file")"
EOF
                    ((successful_digests++))
                    log "✅ $image -> $digest"
                else
                    # Add entry with warning for failed digest fetch
                    cat >> "$LOCK_FILE" << EOF
  "$image":
    digest: "FETCH_FAILED"
    pinned_reference: "$image"
    last_updated: "$(date -Iseconds)"
    source_stack: "$(basename "$stack_file")"
    warning: "Could not fetch digest from registry"
EOF
                    log "❌ Failed to get digest for $image"
                fi
            done <<< "$images"
        fi
    done <<< "$stack_files"
    
    # Add summary to lock file
    cat >> "$LOCK_FILE" << EOF

# Summary
total_images: $total_images
successful_digests: $successful_digests
failed_digests: $((total_images - successful_digests))
EOF
    
    log "✅ Digest lock generation complete"
    log "📊 Total images: $total_images, Successful: $successful_digests, Failed: $((total_images - successful_digests))"
}

# Function to update stack files with pinned digests
update_stacks_with_digests() {
    log "Updating stack files with pinned digests"
    
    if [[ ! -f "$LOCK_FILE" ]]; then
        log "❌ Lock file not found: $LOCK_FILE"
        return 1
    fi
    
    # Create backup directory
    local backup_dir="$PROJECT_ROOT/backups/stacks-$(date +%Y%m%d-%H%M%S)"
    mkdir -p "$backup_dir"
    
    # Process each stack file
    find "$STACKS_DIR" -name "*.yml" -o -name "*.yaml" | while IFS= read -r stack_file; do
        log "Updating $stack_file"
        
        # Create backup
        cp "$stack_file" "$backup_dir/"
        
        # Extract images and update with digests using Python script
        python3 << 'PYTHON_SCRIPT'
import yaml
import sys
import os
import re

stack_file = sys.argv[1] if len(sys.argv) > 1 else ""
lock_file = os.environ.get('LOCK_FILE', '')

if not stack_file or not lock_file or not os.path.exists(lock_file):
    print("Missing required files")
    sys.exit(1)

try:
    # Load lock file
    with open(lock_file, 'r') as f:
        lock_data = yaml.safe_load(f)
    
    # Load stack file
    with open(stack_file, 'r') as f:
        stack_data = yaml.safe_load(f)
    
    # Update images with digests
    if 'services' in stack_data:
        for service_name, service_config in stack_data['services'].items():
            if 'image' in service_config:
                image = service_config['image']
                if image in lock_data.get('images', {}):
                    digest_info = lock_data['images'][image]
                    if digest_info.get('digest') != 'FETCH_FAILED':
                        service_config['image'] = digest_info['pinned_reference']
                        print(f"Updated {service_name}: {image} -> {digest_info['pinned_reference']}")
    
    # Write updated stack file
    with open(stack_file, 'w') as f:
        yaml.dump(stack_data, f, default_flow_style=False, indent=2)
    
except Exception as e:
    print(f"Error processing {stack_file}: {e}")
    sys.exit(1)
PYTHON_SCRIPT "$stack_file"
    done
    
    log "✅ Stack files updated with pinned digests"
    log "📁 Backups stored in: $backup_dir"
}

# Function to validate updated stacks
validate_stacks() {
    log "Validating updated stack files"
    
    local validation_errors=0
    
    find "$STACKS_DIR" -name "*.yml" -o -name "*.yaml" | while IFS= read -r stack_file; do
        # Check YAML syntax
        if ! python3 -c "import yaml; yaml.safe_load(open('$stack_file'))" >/dev/null 2>&1; then
            log "❌ YAML syntax error in $stack_file"
            ((validation_errors++))
        fi
        
        # Check for digest references
        if grep -q '@sha256:' "$stack_file"; then
            log "✅ $stack_file contains digest references"
        else
            log "⚠️  $stack_file does not contain digest references"
        fi
    done
    
    if [[ $validation_errors -eq 0 ]]; then
        log "✅ All stack files validated successfully"
    else
        log "❌ Validation completed with $validation_errors errors"
        return 1
    fi
}

# Function to create cron job for automation
setup_automation() {
    local cron_schedule="0 2 * * 0"  # Weekly on Sunday at 2 AM
    local cron_command="$SCRIPT_DIR/automated-image-update.sh --auto-update"
    
    # Check if cron job already exists
    if crontab -l 2>/dev/null | grep -q "automated-image-update.sh"; then
        log "Cron job already exists for automated image updates"
    else
        # Add cron job
        (crontab -l 2>/dev/null; echo "$cron_schedule $cron_command") | crontab -
        log "✅ Automated weekly image digest updates scheduled"
    fi
}

# Main execution
main() {
    case "${1:-}" in
        "--generate-lock")
            generate_digest_lock
            ;;
        "--update-stacks")
            update_stacks_with_digests
            validate_stacks
            ;;
        "--auto-update")
            generate_digest_lock
            update_stacks_with_digests
            validate_stacks
            ;;
        "--setup-automation")
            setup_automation
            ;;
        "--help"|"-h"|"")
            cat << 'EOF'
Automated Image Digest Management Script

USAGE:
  automated-image-update.sh [OPTIONS]

OPTIONS:
  --generate-lock     Generate digest lock file only
  --update-stacks     Update stack files with pinned digests
  --auto-update       Generate lock and update stacks (full automation)
  --setup-automation  Set up weekly cron job for automated updates
  --help, -h          Show this help message

EXAMPLES:
  # Generate digest lock file
  ./automated-image-update.sh --generate-lock
  
  # Update stack files with digests
  ./automated-image-update.sh --update-stacks
  
  # Full automated update (recommended)
  ./automated-image-update.sh --auto-update
  
  # Set up weekly automation
  ./automated-image-update.sh --setup-automation

NOTES:
  - Requires yq, skopeo, or Docker CLI for fetching digests
  - Creates backups before modifying stack files
  - Logs all operations for auditability
  - Safe to run multiple times (idempotent)
EOF
            ;;
        *)
            log "❌ Unknown option: $1"
            log "Use --help for usage information"
            exit 1
            ;;
    esac
}

# Execute main function with all arguments
main "$@"