#!/bin/bash

set -e

# Restic Data Restore Script for HuggingFace Environment
# Simplified version: Focus on core restore functionality, directly use latest snapshot

# Logging functions
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [RESTIC-RESTORE] $*"
}

log_success() {
    echo -e "\033[32m[$(date '+%Y-%m-%d %H:%M:%S')] [SUCCESS] $*\033[0m"
}

log_error() {
    echo -e "\033[31m[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $*\033[0m"
}

log_warn() {
    echo -e "\033[33m[$(date '+%Y-%m-%d %H:%M:%S')] [WARNING] $*\033[0m"
}

log_info() {
    echo -e "\033[36m[$(date '+%Y-%m-%d %H:%M:%S')] [INFO] $*\033[0m"
}

# Display usage instructions
show_usage() {
    cat << EOF
Restic Data Restore Script (Simplified - Only supports latest snapshot)

Usage: $0 [options]

Options:
  -h, --help              Show this help message
  -l, --list              List available snapshots
  -c, --check             Check repository integrity
  -r, --repo PATH         Specify repository path (default: \$RESTIC_REPOSITORY_PATH)
  -p, --password-file FILE Specify password file (default: \$HOME/config/restic-password)
  -t, --target PATH       Specify restore target path (default: /tmp/restic_restore)
  -m, --mode MODE         Restore mode:
                          - safe: Safe mode, restore to specified directory (default)
                          - direct: Direct mode, restore directly to target location
                          - replace: Replace mode, clear target directory before restore
  -i, --include PATTERN   Only restore matching paths (can be used multiple times)
  -e, --exclude PATTERN   Exclude matching paths (can be used multiple times)
  -v, --verbose           Verbose output
  -f, --force             Force execution without confirmation
  --dry-run               Dry run, don't perform actual restore

Note: This version only supports restoring the latest snapshot, simplifying complex snapshot ID handling logic

Examples:
  # List available snapshots
  $0 --list

  # Safe restore latest snapshot to default temp directory
  $0

  # Direct restore to specified location
  $0 --mode direct --target /home/user

  # Full replacement mode (clear target directory first)
  $0 --mode replace --target /home/user

  # Only restore config files
  $0 --include "/home/user/config"

Environment Variables:
  RESTIC_REPOSITORY_PATH     Restic repository path
  RESTIC_PASSWORD            Restic password
EOF
}

# Check if restic is installed
check_restic() {
    if ! command -v restic &> /dev/null; then
        log_error "restic is not installed, please install restic first"
        return 1
    fi
    return 0
}

# Check and set environment variables
setup_environment() {
    # Set default values
    USER_HOME="${HOME:-/home/user}"
    REPO_PATH="${RESTIC_REPOSITORY_PATH:-$USER_HOME/data/restic-repo}"
    PASSWORD_FILE="${PASSWORD_FILE:-$USER_HOME/config/restic-password}"
    TARGET_PATH="${TARGET_PATH:-/tmp/restic_restore_$(date +%s)}"
    RESTORE_MODE="${RESTORE_MODE:-safe}"
    
    # Check repository path (skip local path check for S3/remote repositories)
    if [[ ! "$REPO_PATH" =~ ^(s3:|sftp:|rest:|rclone:) && ! -d "$REPO_PATH" ]]; then
        log_warn "Local Restic repository does not exist: $REPO_PATH"
        log_warn "This might be the first startup, repository not yet initialized"
        log_warn "Please check RESTIC_REPOSITORY_PATH environment variable or use -r option"
        return 1
    fi
    
    # Set password
    if [[ -n "$RESTIC_PASSWORD" ]]; then
        export RESTIC_PASSWORD
    elif [[ -f "$PASSWORD_FILE" ]]; then
        export RESTIC_PASSWORD="$(cat "$PASSWORD_FILE")"
        log_info "Reading password from file: $PASSWORD_FILE"
    else
        log_error "Password not found, please set RESTIC_PASSWORD environment variable or password file: $PASSWORD_FILE"
        return 1
    fi
    
    # Set repository path
    export RESTIC_REPOSITORY="$REPO_PATH"
    
    log_info "Environment configuration:"
    log_info "  Repository path: $REPO_PATH"
    log_info "  Target path: $TARGET_PATH"
    log_info "  Restore mode: $RESTORE_MODE"
    log_info "  Snapshot: latest (fixed to use latest snapshot)"
    
    return 0
}

# Check repository status
check_repository() {
    log_info "Checking repository status..."
    
    if ! restic snapshots > /dev/null 2>&1; then
        log_warn "Cannot access Restic repository, please check repository path and password"
        log_warn "This might be the first startup or repository not yet initialized"
        return 1
    fi
    
    log_success "Repository status normal"
    return 0
}

# List snapshots
list_snapshots() {
    log_info "Available snapshots list:"
    restic snapshots --compact
}

# Perform restore - directly use latest snapshot
perform_restore() {
    local target_path="$1"
    
    log_info "Restoring latest snapshot to: $target_path"
    
    # Simplified content restore logic
    log_info "Analyzing latest snapshot backup paths, preparing for content restore..."
    
    # Get original backup paths from latest snapshot metadata
    local backup_paths=()
    local snapshot_info
    if snapshot_info=$(restic snapshots --json --latest 1 2>/dev/null); then
        local paths_json
        paths_json=$(echo "$snapshot_info" | jq -r '.[0].paths[]?' 2>/dev/null)
        if [[ -n "$paths_json" ]]; then
            while IFS= read -r path; do
                [[ -n "$path" ]] && backup_paths+=("$path")
            done <<< "$paths_json"
            log_info "Detected backup paths: ${backup_paths[*]}"
        fi
    fi
    
    # Build restore parameters
    local snapshot_target="latest"
    if [[ ${#backup_paths[@]} -gt 0 ]]; then
        # Use first backup path as content root directory
        local content_root="${backup_paths[0]}"
        log_info "Using content restore mode, only restore $content_root contents to $target_path"
        
        # Modify snapshot to latest:subfolder format
        snapshot_target="latest:$content_root"
        log_info "Converting to content restore: $snapshot_target"
    else
        log_warn "Cannot determine backup paths, using traditional full path restore"
    fi
    
    # Replace mode: clear target directory
    if [[ "$RESTORE_MODE" == "replace" ]]; then
        log_info "Replace mode: preparing to clear target directory..."
        
        if [[ -d "$target_path" && -n "$(ls -A "$target_path" 2>/dev/null)" ]]; then
            log_warn "Replace mode will clear target directory: $target_path"
            if [[ "$FORCE" != "true" ]]; then
                read -p "Confirm clear target directory and restore? (y/N): " -n 1 -r
                echo
                [[ ! $REPLY =~ ^[Yy]$ ]] && { log_info "User cancelled operation"; return 1; }
            fi
            
            if [[ "$DRY_RUN" == "true" ]]; then
                log_info "Dry run mode: will clear $target_path"
            else
                log_info "Clearing target directory: $target_path"
                # Safety check: prevent deletion of important system directories
                if [[ "$target_path" =~ ^/(bin|sbin|usr|lib|lib64|etc|proc|sys|dev|boot)$ ]]; then
                    log_error "Refusing to clear system directory: $target_path"
                    return 1
                fi
                
                # Clear directory contents but keep directory itself
                find "$target_path" -mindepth 1 -delete 2>/dev/null || {
                    # If find fails, try deleting one by one
                    for item in "$target_path"/{*,.[^.]*}; do
                        [[ -e "$item" ]] && rm -rf "$item"
                    done
                }
                log_success "Target directory cleared: $target_path"
            fi
        else
            log_info "Target directory is empty or does not exist, skipping clear: $target_path"
        fi
    fi
    
    # If direct mode and target path has content, give warning
    if [[ "$RESTORE_MODE" == "direct" && -d "$target_path" && -n "$(ls -A "$target_path" 2>/dev/null)" ]]; then
        log_warn "Target path is not empty, may overwrite existing files: $target_path"
        if [[ "$FORCE" != "true" ]]; then
            read -p "Continue? (y/N): " -n 1 -r
            echo
            [[ ! $REPLY =~ ^[Yy]$ ]] && { log_info "User cancelled operation"; return 1; }
        fi
    fi
    
    # Create target directory
    mkdir -p "$target_path"
    
    # Build restore command arguments
    local args=("restore" "$snapshot_target" "--target" "$target_path")
    
    # Add include/exclude options
    for pattern in "${INCLUDE_PATTERNS[@]}"; do
        args+=("--include" "$pattern")
    done
    
    for pattern in "${EXCLUDE_PATTERNS[@]}"; do
        args+=("--exclude" "$pattern")
    done
    
    [[ "$VERBOSE" == "true" ]] && args+=("--verbose")
    
    log_info "Executing restore command: restic ${args[*]}"
    
    if [[ "$DRY_RUN" == "true" ]]; then
        log_info "Dry run mode, not performing actual restore"
        return 0
    fi
    
    if restic "${args[@]}"; then
        log_success "Restore completed: $target_path"
        
        # Give specific next steps based on mode
        if [[ "$RESTORE_MODE" == "safe" ]]; then
            log_info "Safe mode: Data restored to safe directory, please check and manually copy to final location"
            log_info "Example: rsync -av $target_path/ /desired/location/"
        elif [[ "$RESTORE_MODE" == "direct" ]]; then
            log_info "Direct mode: Data restored directly to specified directory"
        elif [[ "$RESTORE_MODE" == "replace" ]]; then
            log_info "Replace mode: Target directory intelligently cleared and data restored"
        fi
        
        return 0
    else
        log_error "Restore failed"
        return 1
    fi
}

# Main function
main() {
    local action="restore"
    
    # Initialize arrays
    INCLUDE_PATTERNS=()
    EXCLUDE_PATTERNS=()
    
    # Parse command line arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_usage
                exit 0
                ;;
            -l|--list)
                action="list"
                shift
                ;;
            -c|--check)
                action="check"
                shift
                ;;
            -r|--repo)
                RESTIC_REPOSITORY_PATH="$2"
                shift 2
                ;;
            -p|--password-file)
                PASSWORD_FILE="$2"
                shift 2
                ;;
            -t|--target)
                TARGET_PATH="$2"
                shift 2
                ;;
            -m|--mode)
                RESTORE_MODE="$2"
                if [[ "$RESTORE_MODE" != "safe" && "$RESTORE_MODE" != "direct" && "$RESTORE_MODE" != "replace" ]]; then
                    log_error "Invalid restore mode: $RESTORE_MODE (only supports: safe, direct, replace)"
                    exit 1
                fi
                shift 2
                ;;
            -i|--include)
                INCLUDE_PATTERNS+=("$2")
                shift 2
                ;;
            -e|--exclude)
                EXCLUDE_PATTERNS+=("$2")
                shift 2
                ;;
            -v|--verbose)
                VERBOSE="true"
                shift
                ;;
            -f|--force)
                FORCE="true"
                shift
                ;;
            --dry-run)
                DRY_RUN="true"
                shift
                ;;
            -*)
                log_error "Unknown option: $1"
                show_usage
                exit 1
                ;;
            *)
                log_warn "Ignoring argument: $1 (this version only supports latest snapshot)"
                shift
                ;;
        esac
    done
    
    # Check restic
    if ! check_restic; then
        exit 1
    fi
    
    # Setup environment
    if ! setup_environment; then
        exit 1
    fi
    
    # Check repository
    if ! check_repository; then
        exit 1
    fi
    
    # Execute corresponding operation
    case $action in
        list)
            list_snapshots
            ;;
        check)
            log_info "Performing repository integrity check..."
            if restic check; then
                log_success "Repository integrity check passed"
            else
                log_error "Repository integrity check failed"
                exit 1
            fi
            ;;
        restore)
            # Execute restore directly, using latest snapshot
            perform_restore "$TARGET_PATH"
            ;;
    esac
}

# Script entry point
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi