#!/bin/bash
# Docker Stack Backup Script
# Enhanced version matching Go docwell functionality

set -euo pipefail

# Version
VERSION="2.6.2"

# Resolve script directory and source shared library
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib/common.sh
source "$SCRIPT_DIR/lib/common.sh"

# Enable debug trace after sourcing common.sh (so DEBUG variable is available)
enable_debug_trace

# Default configuration
HOSTNAME=$(hostname 2>/dev/null | tr -cd 'a-zA-Z0-9._-' || echo "unknown")
DATE=$(date +%Y-%m-%d)
BACKUP_BASE="${BACKUP_BASE:-/storage/backups/docker-$HOSTNAME}"
BACKUP_DIR="$BACKUP_BASE/$DATE"
LOG_FILE="${LOG_FILE:-/tmp/docwell/docker-backup.log}"

# CLI flags
BACKUP_ALL=false
BACKUP_STACK=""
LIST_ONLY=false
AUTO_MODE=false
COMPRESSION="zstd"

# Parser configuration
MAX_STACK_NAME_LEN=30

# Load config file overrides
load_config

# Parse command line arguments
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --version)
                echo "docker-backup.sh v$VERSION"
                exit 0
                ;;
            --quiet|-q)
                QUIET=true
                shift
                ;;
            --yes|-y)
                YES=true
                shift
                ;;
            --all|-a)
                BACKUP_ALL=true
                shift
                ;;
            --stack|-s)
                require_arg "$1" "${2:-}" || exit 1
                BACKUP_STACK="$2"
                shift 2
                ;;
            --list|-l)
                LIST_ONLY=true
                shift
                ;;
            --stacks-dir)
                STACKS_DIR="$2"
                shift 2
                ;;
            --backup-base)
                BACKUP_BASE="$2"
                BACKUP_DIR="$BACKUP_BASE/$DATE"
                shift 2
                ;;
            --auto)
                AUTO_MODE=true
                QUIET=true
                YES=true
                shift
                ;;
            --install-deps)
                AUTO_INSTALL=true
                shift
                ;;
            --log)
                LOG_FILE="$2"
                shift 2
                ;;
            --compression|-c)
                COMPRESSION="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --debug|-d)
                DEBUG=true
                shift
                ;;
            --verbose|-v)
                VERBOSE=true
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                echo -e "${RED}Unknown option: $1${NC}" >&2
                show_help
                exit 1
                ;;
        esac
    done
}

show_help() {
    cat << EOF
Docker Stack Backup Script v$VERSION

Usage: $0 [OPTIONS]

Options:
    --all, -a              Backup all stacks
    --stack STACK, -s      Backup specific stack by name
    --list, -l             List available stacks for backup
    --stacks-dir DIR       Stacks directory (default: /opt/stacks)
    --backup-base DIR      Backup base directory (default: /storage/backups/docker-\$HOSTNAME)
    --log FILE             Log file path (default: /tmp/docwell/docker-backup.log)
    --quiet, -q            Suppress non-error output
    --yes, -y              Auto-confirm all prompts
    --auto                 Auto mode: backup all stacks & clean resources (for cron)
    --compression METHOD   Compression method: zstd or gzip (default: zstd)
    --dry-run              Show what would be done without actually doing it
    --debug, -d            Enable debug mode (verbose output + xtrace)
    --verbose, -v          Enable verbose logging
    --install-deps         Auto-install missing dependencies (requires root)
    --version              Show version information
    --help, -h             Show this help message

Examples:
    $0 --list                    # List all stacks
    $0 --all                     # Backup all stacks
    $0 --stack myapp             # Backup specific stack
    $0 --all --quiet --yes       # Backup all stacks non-interactively
    $0 --all --dry-run           # Preview backup without executing
EOF
}

# Check for required dependencies
check_dependencies() {
    local required_cmds=("docker")
    local optional_cmds=("zstd")
    local missing=()
    local missing_optional=()

    for cmd in "${required_cmds[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    for cmd in "${optional_cmds[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing_optional+=("$cmd")
        fi
    done

    if [[ ${#missing_optional[@]} -gt 0 ]]; then
        log_warn "Optional dependencies missing: ${missing_optional[*]} (will use fallback)"
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        log_error "Missing required dependencies: ${missing[*]}"
        echo -e "${BLUE}Install commands:${NC}"
        command -v apt &>/dev/null && echo -e "  ${GREEN}Debian/Ubuntu:${NC} sudo apt install docker.io ${missing_optional[*]}"
        command -v pacman &>/dev/null && echo -e "  ${GREEN}Arch Linux:${NC} sudo pacman -S docker ${missing_optional[*]}"

        if [[ "$AUTO_INSTALL" == "true" ]]; then
            log_info "Auto-installing dependencies..."
            install_dependencies "${missing[@]}" "${missing_optional[@]}"
            return $?
        fi
        return 1
    fi
    return 0
}

find_compose_file() {
    local stack="$1"
    get_compose_file "$STACKS_DIR/$stack"
}

stop_stack() {
    local stack="$1"
    local stack_path="$STACKS_DIR/$stack"
    local compose_file

    log_debug "Attempting to stop stack: $stack"
    compose_file=$(find_compose_file "$stack") || {
        log_debug "No compose file found for stack: $stack"
        return 1
    }

    if is_stack_running "$stack"; then
        log_info "Stopping $stack..."
        if run_or_dry "stop $stack" docker compose -f "$stack_path/$compose_file" down >/dev/null 2>&1; then
            log_debug "Stack stopped successfully: $stack"
            return 0
        else
            log_error "Failed to stop stack: $stack"
            return 1
        fi
    fi
    log_debug "Stack was not running: $stack"
    return 1
}

start_stack() {
    local stack="$1"
    local stack_path="$STACKS_DIR/$stack"
    local compose_file

    log_debug "Attempting to start stack: $stack"
    compose_file=$(find_compose_file "$stack") || {
        log_error "No compose file found for stack: $stack"
        return 1
    }

    log_info "Starting $stack..."
    if run_or_dry "start $stack" docker compose -f "$stack_path/$compose_file" up -d >/dev/null 2>&1; then
        log_debug "Stack started successfully: $stack"
        return 0
    else
        log_error "Failed to start stack: $stack"
        return 1
    fi
}

create_backup_archive() {
    local stack="$1"
    log_debug "Creating backup archive for: $stack"

    # Temp files in BACKUP_DIR, cleaned up on failure
    local temp_tar
    temp_tar=$(mktemp "${BACKUP_DIR}/.${stack}.tar.XXXXXX" 2>/dev/null) || {
        log_error "Failed to create temporary file in $BACKUP_DIR"
        return 1
    }
    register_cleanup "rm -f '$temp_tar'"
    log_debug "Temporary tar file: $temp_tar"

    log_info "Archiving $stack..."

    local comp_info
    comp_info=$(get_compressor "$COMPRESSION")
    local compressor="${comp_info%%:*}"
    local ext="${comp_info#*:}"
    log_debug "Using compressor: $compressor, extension: $ext"

    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "[DRY-RUN] Would create: $BACKUP_DIR/docker-$stack$ext"
        rm -f "$temp_tar"
        return 0
    fi

    log_debug "Creating tar archive..."
    if ! tar -cf "$temp_tar" -C "$STACKS_DIR" "$stack" 2>/dev/null; then
        log_error "Failed to create tar archive for $stack"
        rm -f "$temp_tar"
        return 1
    fi
    
    local tar_size
    tar_size=$(stat -f%z "$temp_tar" 2>/dev/null || stat -c%s "$temp_tar" 2>/dev/null || echo "unknown")
    log_debug "Tar archive size: $tar_size bytes"
    
    log_debug "Compressing archive with $compressor..."
    if ! $compressor < "$temp_tar" > "$BACKUP_DIR/docker-$stack$ext" 2>/dev/null; then
        log_error "Failed to compress archive for $stack"
        rm -f "$temp_tar" "$BACKUP_DIR/docker-$stack$ext" 2>/dev/null
        return 1
    fi
    
    local final_size
    final_size=$(stat -f%z "$BACKUP_DIR/docker-$stack$ext" 2>/dev/null || stat -c%s "$BACKUP_DIR/docker-$stack$ext" 2>/dev/null || echo "unknown")
    log_debug "Final backup size: $final_size bytes"
    log_debug "Backup archive created successfully: $BACKUP_DIR/docker-$stack$ext"
    
    rm -f "$temp_tar"
    return 0
}

backup_stack() {
    local stack="$1"
    log_debug "Starting backup for stack: $stack"

    if ! validate_stack_name "$stack"; then
        log_error "Invalid stack name: $stack"
        return 1
    fi

    local stack_path="$STACKS_DIR/$stack"
    if [[ ! -d "$stack_path" ]]; then
        log_error "Stack directory not found: $stack_path"
        return 1
    fi

    log_debug "Stack path verified: $stack_path"

    # Create backup directory
    if ! mkdir -p "$BACKUP_DIR"; then
        log_error "Failed to create backup directory: $BACKUP_DIR"
        return 1
    fi
    log_debug "Backup directory ready: $BACKUP_DIR"

    local was_running=false

    if stop_stack "$stack"; then
        was_running=true
        log_debug "Stack was running, stopped for backup"
    else
        log_debug "Stack was not running"
    fi

    if create_backup_archive "$stack"; then
        log_success "Backup complete: $stack"
        if [[ "$was_running" == "true" ]]; then
            log_debug "Restarting stack: $stack"
            start_stack "$stack" || log_warn "Failed to restart $stack"
        fi
        return 0
    else
        log_error "Backup failed: $stack"
        if [[ "$was_running" == "true" ]]; then
            log_debug "Attempting to restart stack after failed backup: $stack"
            start_stack "$stack" || log_warn "Failed to restart $stack"
        fi
        return 1
    fi
}

list_stacks() {
    show_spinner "Loading stacks..."
    local stacks
    mapfile -t stacks < <(get_stacks)
    hide_spinner

    if [[ ${#stacks[@]} -eq 0 ]]; then
        log_info "No stacks found in $STACKS_DIR"
        return 1
    fi

    for stack in "${stacks[@]}"; do
        local status="stopped"
        local size
        size=$(get_stack_size "$stack")
        if is_stack_running "$stack"; then
            status="running"
        fi
        echo -e "$stack\t$status\t$size"
    done
}

interactive_backup() {
    local stacks
    show_spinner "Loading stacks..."
    mapfile -t stacks < <(get_stacks)
    hide_spinner

    if [[ ${#stacks[@]} -eq 0 ]]; then
        log_error "No stacks found in $STACKS_DIR"
        exit 1
    fi

    log_header "Backup Docker Stacks"
    echo -e "${GRAY}DocWell Backup v$VERSION${NC}"
    echo

    echo "Available stacks:"
    echo -e "  ${BOLD}0${NC}) All stacks ${GRAY}[Backup everything]${NC}"
    for i in "${!stacks[@]}"; do
        local stack="${stacks[$i]}"
        local status="●"
        local color="$GRAY"
        if is_stack_running "$stack"; then
            color="$GREEN"
        fi
        local size
        size=$(get_stack_size "$stack")
        echo -e "  ${BOLD}$((i+1))${NC}) ${color}${status}${NC} ${stack:0:$MAX_STACK_NAME_LEN} ${GRAY}[${size}]${NC}"
    done
    echo

    read -rp "Enter selection (0 for all, comma-separated numbers): " selection

    log_info "Starting backup on $HOSTNAME"
    local start_time=$SECONDS

    if [[ "$selection" == "0" ]]; then
        _backup_stacks_parallel "${stacks[@]}"
    else
        IFS=',' read -ra SELECTIONS <<< "$selection"
        for sel in "${SELECTIONS[@]}"; do
            sel=$(echo "$sel" | xargs)

            if [[ ! "$sel" =~ ^[0-9]+$ ]]; then
                 log_warn "Invalid selection: $sel"
                 continue
            fi

            if [[ "$sel" -ge 1 ]] && [[ "$sel" -le "${#stacks[@]}" ]]; then
                local idx=$((sel-1))
                backup_stack "${stacks[$idx]}" || true
            else
                log_warn "Selection out of range: $sel"
            fi
        done
    fi

    local elapsed=$(( SECONDS - start_time ))
    log_info "Backup completed in $(format_elapsed "$elapsed")"
}

_backup_stacks_parallel() {
    local stacks=("$@")
    local total=${#stacks[@]}
    local pids=()
    local fail_dir
    fail_dir=$(mktemp -d /tmp/backup_fail.XXXXXX)
    register_cleanup "rm -rf '$fail_dir'"

    log_info "Starting backups for $total stack(s)..."

    local idx=0
    for stack in "${stacks[@]}"; do
        ((idx++))
        echo -e "${CYAN}[$idx/$total]${NC} Backing up $stack..."
        parallel_throttle "$MAX_PARALLEL"
        (
            if ! backup_stack "$stack"; then
                touch "$fail_dir/$stack"
            fi
        ) &
        _parallel_pids+=($!)
    done

    parallel_wait

    # Collect failures
    local failed_stacks=()
    for stack in "${stacks[@]}"; do
        if [[ -f "$fail_dir/$stack" ]]; then
            failed_stacks+=("$stack")
        fi
    done

    if [[ ${#failed_stacks[@]} -gt 0 ]]; then
        log_warn "Failed to backup ${#failed_stacks[@]} stack(s): ${failed_stacks[*]}"
    else
        log_success "All stacks backed up successfully"
    fi
}

auto_mode() {
    log_info "Starting auto mode on $HOSTNAME"
    local start_time=$SECONDS

    # Backup all stacks
    local stacks
    mapfile -t stacks < <(get_stacks)

    if [[ ${#stacks[@]} -eq 0 ]]; then
        log_warn "No stacks found for backup"
    else
        log_info "Backing up ${#stacks[@]} stack(s)..."
        local failed_count=0
        local idx=0
        for stack in "${stacks[@]}"; do
            ((idx++))
            [[ "$QUIET" == "false" ]] && echo -e "${CYAN}[$idx/${#stacks[@]}]${NC} $stack"
            if ! backup_stack "$stack"; then
                ((failed_count++))
            fi
        done

        if [[ $failed_count -gt 0 ]]; then
            log_warn "Failed to backup $failed_count stack(s)"
        else
            log_success "All stacks backed up successfully"
        fi
    fi

    # Cleanup: stopped containers
    log_info "Cleaning up stopped containers..."
    if run_or_dry "prune containers" docker container prune -f > /dev/null 2>&1; then
        log_success "Stopped containers cleaned"
    else
        log_warn "Failed to clean stopped containers"
    fi

    # Cleanup: dangling images
    log_info "Cleaning up dangling images..."
    if run_or_dry "prune images" docker image prune -f > /dev/null 2>&1; then
        log_success "Dangling images cleaned"
    else
        log_warn "Failed to clean dangling images"
    fi

    local elapsed=$(( SECONDS - start_time ))
    log_info "Auto mode completed in $(format_elapsed "$elapsed")"
}

main() {
    parse_args "$@"
    
    log_debug "Starting docker-backup.sh v$VERSION"
    log_debug "Configuration: STACKS_DIR=$STACKS_DIR, BACKUP_BASE=$BACKUP_BASE"
    log_debug "Flags: QUIET=$QUIET, YES=$YES, DRY_RUN=$DRY_RUN, DEBUG=$DEBUG, VERBOSE=$VERBOSE"

    # Check dependencies first
    if ! check_dependencies; then
        log_error "Dependency check failed"
        exit 1
    fi

    # Validation of critical dirs
    if [[ ! -d "$STACKS_DIR" ]]; then
        log_error "Stacks directory does not exist: $STACKS_DIR"
        log_debug "Attempting to create stacks directory..."
        if mkdir -p "$STACKS_DIR" 2>/dev/null; then
            log_info "Created stacks directory: $STACKS_DIR"
        else
            log_error "Cannot create stacks directory: $STACKS_DIR"
            exit 1
        fi
    fi

    # Check docker connectivity
    if ! docker info >/dev/null 2>&1; then
        log_debug "Docker info check failed, checking permissions..."
        if [[ "$EUID" -ne 0 ]]; then
             if ! confirm "Docker socket not accessible. Run with sudo?"; then
                exit 1
             fi
             log_debug "Re-executing with sudo..."
             exec sudo "$SCRIPT_DIR/$(basename "${BASH_SOURCE[0]}")" "$@"
        else
            log_error "Cannot connect to Docker daemon even as root."
            exit 1
        fi
    fi

    # Handle auto mode
    if [[ "$AUTO_MODE" == "true" ]]; then
        auto_mode
        exit 0
    fi

    # Handle list only
    if [[ "$LIST_ONLY" == "true" ]]; then
        list_stacks
        exit 0
    fi

    # Handle specific stack backup
    if [[ -n "$BACKUP_STACK" ]]; then
        if ! validate_stack_name "$BACKUP_STACK"; then
            exit 1
        fi
        local start_time=$SECONDS
        backup_stack "$BACKUP_STACK"
        local rc=$?
        local elapsed=$(( SECONDS - start_time ))
        log_info "Completed in $(format_elapsed "$elapsed")"
        exit $rc
    fi

    # Handle backup all
    if [[ "$BACKUP_ALL" == "true" ]]; then
        local stacks
        mapfile -t stacks < <(get_stacks)
        if [[ ${#stacks[@]} -eq 0 ]]; then
            log_error "No stacks found"
            exit 1
        fi

        local start_time=$SECONDS
        _backup_stacks_parallel "${stacks[@]}"
        local elapsed=$(( SECONDS - start_time ))
        log_info "Completed in $(format_elapsed "$elapsed")"
        exit 0
    fi

    # Interactive mode
    interactive_backup
}

main "$@"