#!/bin/bash

# Exit on error
set -e

# Load environment variables
if [ -f .env ]; then
    source .env
fi

# Default values
BASE_DIR=${COMPOSESYNC_BASE_DIR:-"/opt/composesync/stacks"}
KEEP_VERSIONS=${KEEP_VERSIONS:-10}
UPDATE_INTERVAL=${UPDATE_INTERVAL_SECONDS:-3600}
UPDATE_MODE=${UPDATE_MODE:-"notify_and_apply"}
DRY_RUN=${DRY_RUN:-"false"}
STACKS=${STACKS:-1}

# Function to log messages
log() {
    local prefix=""
    if [ "$DRY_RUN" = "true" ]; then
        prefix="[DRY-RUN] "
    fi
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ${prefix}$1"
}

# Function to acquire lock
acquire_lock() {
    local lock_file="$1"
    local lock_timeout=300  # 5 minutes timeout
    
    # Try to create the lock file
    if mkdir "$lock_file" 2>/dev/null; then
        # Successfully acquired lock
        return 0
    else
        # Check if lock is stale
        if [ -d "$lock_file" ]; then
            local lock_age=$(($(date +%s) - $(stat -c %Y "$lock_file")))
            if [ $lock_age -gt $lock_timeout ]; then
                log "Found stale lock, removing..."
                rm -rf "$lock_file"
                if mkdir "$lock_file" 2>/dev/null; then
                    return 0
                fi
            fi
        fi
        return 1
    fi
}

# Function to release lock
release_lock() {
    local lock_file="$1"
    rm -rf "$lock_file"
}

# Function to download a file
download_file() {
    local url=$1
    local output=$2
    local tool=$3
    local subpath=$4
    local git_ref=$5

    log "Downloading $url to $output"
    
    case $tool in
        "wget")
            if ! wget -q -O "$output" "$url"; then
                log "ERROR: Failed to download $url"
                return 1
            fi
            # Verify file is not empty
            if [ ! -s "$output" ]; then
                log "ERROR: Downloaded file $output is empty"
                return 1
            fi
            ;;
        "git")
            if [ ! -d "${output}.git" ]; then
                if ! git clone --quiet "$url" "${output}.git"; then
                    log "ERROR: Failed to clone $url"
                    return 1
                fi
            else
                if ! (cd "${output}.git" && git fetch --quiet); then
                    log "ERROR: Failed to fetch from $url"
                    return 1
                fi
            fi
            
            # Checkout specific branch/tag if specified
            if [ -n "$git_ref" ]; then
                if ! (cd "${output}.git" && git checkout --quiet "$git_ref"); then
                    log "ERROR: Failed to checkout $git_ref"
                    return 1
                fi
            fi
            
            if [ -n "$subpath" ]; then
                if [ ! -f "${output}.git/$subpath" ]; then
                    log "ERROR: Subpath $subpath not found in repository"
                    return 1
                fi
                cp "${output}.git/$subpath" "$output"
            else
                if [ ! -f "${output}.git/docker-compose.yml" ]; then
                    log "ERROR: docker-compose.yml not found in repository root"
                    return 1
                fi
                cp "${output}.git/docker-compose.yml" "$output"
            fi
            ;;
        *)
            log "ERROR: Unsupported tool $tool"
            return 1
            ;;
    esac
}

# Function to get version identifier
get_version_id() {
    local path=$1
    local tool=$2

    case $tool in
        "git")
            if [ -d "${path}.git" ]; then
                git -C "${path}.git" rev-parse --short HEAD
            else
                date +%Y%m%d%H%M%S
            fi
            ;;
        *)
            date +%Y%m%d%H%M%S
            ;;
    esac
}

# Function to clean up old backup directories
cleanup_old_backups() {
    local path=$1
    local keep_versions=$2
    
    log "Cleaning up old backup directories (keeping $keep_versions)"
    find "$path/backups" -maxdepth 1 -type d -name "backup-*" | sort -r | tail -n +$((keep_versions + 1)) | xargs -r rm -rf
}

# Function to send webhook notification
send_webhook_notification() {
    local event="$1"
    local stack_name="$2"
    local message="$3"
    local version_id="$4"
    local diff="$5"

    if [ -z "$NOTIFICATION_WEBHOOK_URL" ]; then
        return
    fi

    # Escape newlines and quotes in diff for JSON
    local diff_json
    if [ -n "$diff" ]; then
        diff_json=$(echo "$diff" | head -50 | sed 's/\\/\\\\/g; s/\"/\\\"/g; s/$/\\n/' | tr -d '\r' | tr -d '\000')
    else
        diff_json=""
    fi

    local payload="{\n"
    payload+="  \"event\": \"$event\",\n"
    payload+="  \"stack_name\": \"$stack_name\",\n"
    payload+="  \"timestamp\": \"$(date -u +%Y-%m-%dT%H:%M:%SZ)\",\n"
    payload+="  \"message\": \"$message\",\n"
    payload+="  \"version_id\": \"$version_id\",\n"
    payload+="  \"diff\": \"$diff_json\"\n"
    payload+="}"

    curl -s -X POST -H "Content-Type: application/json" -d "$payload" "$NOTIFICATION_WEBHOOK_URL" >/dev/null 2>&1 || true
}

# Function to process a stack
process_stack() {
    local stack_num=$1
    local name_var="STACK_${stack_num}_NAME"
    local url_var="STACK_${stack_num}_URL"
    local path_var="STACK_${stack_num}_PATH"
    local tool_var="STACK_${stack_num}_TOOL"
    local interval_var="STACK_${stack_num}_INTERVAL"
    local keep_versions_var="STACK_${stack_num}_KEEP_VERSIONS"
    local git_subpath_var="STACK_${stack_num}_GIT_SUBPATH"
    local git_ref_var="STACK_${stack_num}_GIT_REF"

    local name=${!name_var}
    local url=${!url_var}
    local path=${!path_var}
    local tool=${!tool_var}
    local interval=${!interval_var:-$UPDATE_INTERVAL}
    local keep_versions=${!keep_versions_var:-$KEEP_VERSIONS}
    local git_subpath=${!git_subpath_var}
    local git_ref=${!git_ref_var}

    if [ -z "$name" ] || [ -z "$url" ] || [ -z "$path" ] || [ -z "$tool" ]; then
        log "Error: Missing required configuration for stack $stack_num"
        return 1
    fi

    # Create lock file path
    local lock_file="$path/.lock"
    
    # Try to acquire lock
    if ! acquire_lock "$lock_file"; then
        log "Warning: Could not acquire lock for stack $name, skipping..."
        return 1
    fi
    
    # Ensure lock is released on exit
    trap 'release_lock "$lock_file"' EXIT

    log "Processing stack $name"

    # Create stack directory and backups directory if they don't exist
    mkdir -p "$path"
    mkdir -p "$path/backups"

    # Download main compose file
    local compose_file="$path/docker-compose.yml"
    if ! download_file "$url" "$compose_file" "$tool" "$git_subpath" "$git_ref"; then
        log "ERROR: Failed to download main compose file for stack $name, skipping..."
        return 1
    fi

    # Get version identifier
    local version_id=$(get_version_id "$path" "$tool")
    log "Version ID: $version_id"

    # Create backup directory for this update
    local backup_dir="$path/backups/backup-$(date +%Y%m%d%H%M%S)"
    if [ "$DRY_RUN" != "true" ]; then
        mkdir -p "$backup_dir"
    fi

    # Download extra files if configured
    # Support custom ordering via STACK_N_EXTRA_FILES_ORDER (comma-separated list of numbers)
    local extra_file_order_var="STACK_${stack_num}_EXTRA_FILES_ORDER"
    local extra_file_order=${!extra_file_order_var}
    local extra_files=()
    if [ -n "$extra_file_order" ]; then
        # Custom order specified
        IFS=',' read -ra order_arr <<< "$extra_file_order"
        for idx in "${order_arr[@]}"; do
            idx=$(echo "$idx" | xargs) # trim whitespace
            local extra_file_var="STACK_${stack_num}_EXTRA_FILES_${idx}"
            local extra_file_url=${!extra_file_var}
            if [ -n "$extra_file_url" ]; then
                local extra_filename=$(basename "$extra_file_url")
                local extra_file_path="$path/$extra_filename"
                log "Downloading extra file $idx (custom order): $extra_filename"
                if ! download_file "$extra_file_url" "$extra_file_path" "wget"; then
                    log "ERROR: Failed to download extra file $extra_filename, skipping..."
                    continue
                fi
                extra_files+=("$extra_file_path")
                # Version the extra file
                local versioned_extra_file="$path/compose-${version_id}-${extra_filename}.bak"
                if [ "$DRY_RUN" != "true" ]; then
                    cp "$extra_file_path" "$versioned_extra_file"
                    cp "$extra_file_path" "$backup_dir/$extra_filename"
                else
                    log "Would create versioned file: $versioned_extra_file"
                    log "Would backup file to: $backup_dir/$extra_filename"
                fi
            fi
        done
    else
        # Default: process in numeric order starting from 1
        local extra_file_num=1
        while true; do
            local extra_file_var="STACK_${stack_num}_EXTRA_FILES_${extra_file_num}"
            local extra_file_url=${!extra_file_var}
            if [ -z "$extra_file_url" ]; then
                break
            fi
            local extra_filename=$(basename "$extra_file_url")
            local extra_file_path="$path/$extra_filename"
            log "Downloading extra file $extra_file_num: $extra_filename"
            if ! download_file "$extra_file_url" "$extra_file_path" "wget"; then
                log "ERROR: Failed to download extra file $extra_filename, skipping..."
                ((extra_file_num++))
                continue
            fi
            extra_files+=("$extra_file_path")
            # Version the extra file
            local versioned_extra_file="$path/compose-${version_id}-${extra_filename}.bak"
            if [ "$DRY_RUN" != "true" ]; then
                cp "$extra_file_path" "$versioned_extra_file"
                cp "$extra_file_path" "$backup_dir/$extra_filename"
            else
                log "Would create versioned file: $versioned_extra_file"
                log "Would backup file to: $backup_dir/$extra_filename"
            fi
            ((extra_file_num++))
        done
    fi

    # Create versioned copy of main compose file
    local versioned_file="$path/compose-${version_id}.yml.bak"
    if [ "$DRY_RUN" != "true" ]; then
        cp "$compose_file" "$versioned_file"
        # Backup the main compose file
        cp "$compose_file" "$backup_dir/docker-compose.yml"
        # Backup the override file if it exists
        if [ -f "$path/docker-compose.override.yml" ]; then
            cp "$path/docker-compose.override.yml" "$backup_dir/docker-compose.override.yml"
        fi
    else
        log "Would create versioned file: $versioned_file"
        log "Would backup main compose file to: $backup_dir/docker-compose.yml"
        if [ -f "$path/docker-compose.override.yml" ]; then
            log "Would backup override file to: $backup_dir/docker-compose.override.yml"
        fi
    fi

    # Clean up old versions and backups
    if [ "$DRY_RUN" != "true" ]; then
        log "Cleaning up old versions (keeping $keep_versions)"
        find "$path" -name "compose-*.yml.bak" -type f | sort -r | tail -n +$((keep_versions + 1)) | xargs -r rm
        cleanup_old_backups "$path" "$keep_versions"
    else
        log "Would clean up old versions (keeping $keep_versions)"
        log "Would clean up old backup directories (keeping $keep_versions)"
    fi

    # Check for changes
    if [ -f "$compose_file" ] && [ -f "$versioned_file" ]; then
        if ! diff -q "$compose_file" "$versioned_file" >/dev/null; then
            log "Changes detected in $name"
            # Generate diff (truncated to 50 lines)
            local diff_output
            diff_output=$(diff -u "$versioned_file" "$compose_file" | head -50)
            if [ $(diff -u "$versioned_file" "$compose_file" | wc -l) -gt 50 ]; then
                diff_output="$diff_output\n(diff truncated, showing first 50 lines)"
            fi
            if [ "$UPDATE_MODE" = "notify_and_apply" ]; then
                if [ "$DRY_RUN" = "true" ]; then
                    log "DRY-RUN: Would apply changes to $name"
                    
                    # Build docker compose command with all files (safer array-based approach)
                    local compose_cmd_array=("docker" "compose" "-f" "$compose_file")
                    for extra_file in "${extra_files[@]}"; do
                        compose_cmd_array+=("-f" "$extra_file")
                    done
                    if [ -f "$path/docker-compose.override.yml" ]; then
                        compose_cmd_array+=("-f" "$path/docker-compose.override.yml")
                    fi
                    compose_cmd_array+=("up" "-d")
                    
                    log "DRY-RUN: Would run: ${compose_cmd_array[*]}"
                    log "DRY-RUN: Changes that would be applied:"
                    echo "$diff_output"
                    send_webhook_notification "update_success" "$name" "[DRY-RUN] Would apply changes to $name" "$version_id" "$diff_output"
                else
                    log "Applying changes to $name"
                    
                    # Build docker compose command with all files (safer array-based approach)
                    local compose_cmd_array=("docker" "compose" "-f" "$compose_file")
                    for extra_file in "${extra_files[@]}"; do
                        compose_cmd_array+=("-f" "$extra_file")
                    done
                    if [ -f "$path/docker-compose.override.yml" ]; then
                        compose_cmd_array+=("-f" "$path/docker-compose.override.yml")
                    fi
                    compose_cmd_array+=("up" "-d")
                    
                    log "Running: ${compose_cmd_array[*]}"
                    if "${compose_cmd_array[@]}"; then
                        log "Successfully updated stack $name"
                        send_webhook_notification "update_success" "$name" "Successfully updated stack $name" "$version_id" "$diff_output"
                    else
                        log "ERROR: Failed to update stack $name, attempting rollback..."
                        send_webhook_notification "update_failure" "$name" "Failed to update stack $name, rolled back to previous version" "$version_id" "$diff_output"
                        
                        # Rollback: restore from backup
                        if [ -f "$backup_dir/docker-compose.yml" ]; then
                            cp "$backup_dir/docker-compose.yml" "$compose_file"
                            log "Restored main compose file from backup"
                        fi
                        
                        # Restore extra files
                        for extra_file in "${extra_files[@]}"; do
                            local extra_filename=$(basename "$extra_file")
                            if [ -f "$backup_dir/$extra_filename" ]; then
                                cp "$backup_dir/$extra_filename" "$extra_file"
                                log "Restored $extra_filename from backup"
                            fi
                        done
                        
                        # Restore override file if it existed
                        if [ -f "$backup_dir/docker-compose.override.yml" ]; then
                            cp "$backup_dir/docker-compose.override.yml" "$path/docker-compose.override.yml"
                            log "Restored override file from backup"
                        fi
                        
                        # Try to restart with rolled back configuration
                        log "Attempting to restart stack with rolled back configuration..."
                        if "${compose_cmd_array[@]}"; then
                            log "Successfully rolled back stack $name"
                        else
                            log "CRITICAL: Failed to rollback stack $name - manual intervention required"
                        fi
                    fi
                fi
            else
                log "Changes detected but not applied (notify_only mode)"
                send_webhook_notification "update_success" "$name" "Changes detected in $name (notify_only mode)" "$version_id" "$diff_output"
            fi
        else
            log "No changes detected in $name"
        fi
    fi
}

# Main loop
while true; do
    log "Starting update check"
    
    for ((i=1; i<=STACKS; i++)); do
        process_stack $i
    done
    
    log "Update check complete. Sleeping for $UPDATE_INTERVAL seconds"
    sleep $UPDATE_INTERVAL
done