#!/bin/bash

# Immich Restore Script
# This script restores an Immich installation from backups created by backup-immich.sh
# Based on recommendations from https://immich.app/docs/administration/backup-and-restore/

# Set up error handling
set -e

# Load environment variables from the .env file
ENV_FILE="$(dirname "$0")/../.env"
if [ -f "$ENV_FILE" ]; then
    echo "Loading environment variables from $ENV_FILE"
    source "$ENV_FILE"
else
    echo "Error: .env file not found in $(dirname "$0")/.."
    exit 1
fi

# Set up logging to central logs directory
LOG_DIR="$(dirname "$0")/../logs"
mkdir -p "$LOG_DIR"
LOG_FILE="${LOG_DIR}/immich-restore.log"

# Function to log with timestamp
log_message() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$LOG_FILE"
}

# Function to display usage
usage() {
    echo "Usage: $0 --db-backup <path> --uploads-backup <path> [options]"
    echo ""
    echo "Required arguments:"
    echo "  --db-backup PATH        Path to database backup file (.sql.gz)"
    echo "  --uploads-backup PATH   Path to uploads backup file (.tar.gz)"
    echo ""
    echo "Optional arguments:"
    echo "  --dry-run              Show what would be restored without making changes"
    echo "  --skip-db              Skip database restoration"
    echo "  --skip-uploads         Skip uploads restoration"
    echo "  --help                 Show this help message"
    echo ""
    echo "Example:"
    echo "  $0 --db-backup ./immich_backups/immich_db_backup_20250526_120000.sql.gz \\"
    echo "     --uploads-backup ./immich_backups/immich_uploads_20250526_120000.tar.gz"
}

# Parse command line arguments
DB_BACKUP=""
UPLOADS_BACKUP=""
DRY_RUN=false
SKIP_DB=false
SKIP_UPLOADS=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --db-backup)
            DB_BACKUP="$2"
            shift 2
            ;;
        --uploads-backup)
            UPLOADS_BACKUP="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        --skip-db)
            SKIP_DB=true
            shift
            ;;
        --skip-uploads)
            SKIP_UPLOADS=true
            shift
            ;;
        --help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Validate required arguments
if [ -z "$DB_BACKUP" ] && [ "$SKIP_DB" = false ]; then
    echo "Error: --db-backup is required unless --skip-db is specified"
    usage
    exit 1
fi

if [ -z "$UPLOADS_BACKUP" ] && [ "$SKIP_UPLOADS" = false ]; then
    echo "Error: --uploads-backup is required unless --skip-uploads is specified"
    usage
    exit 1
fi

# Validate backup files exist
if [ "$SKIP_DB" = false ] && [ ! -f "$DB_BACKUP" ]; then
    echo "Error: Database backup file not found: $DB_BACKUP"
    exit 1
fi

if [ "$SKIP_UPLOADS" = false ] && [ ! -f "$UPLOADS_BACKUP" ]; then
    echo "Error: Uploads backup file not found: $UPLOADS_BACKUP"
    exit 1
fi

echo "=== IMMICH RESTORE OPERATION ==="
echo "Database backup: ${DB_BACKUP:-SKIPPED}"
echo "Uploads backup: ${UPLOADS_BACKUP:-SKIPPED}"
echo "Dry run mode: $DRY_RUN"
echo ""

if [ "$DRY_RUN" = true ]; then
    echo "DRY RUN MODE - No changes will be made"
    echo ""
fi

# Verify required environment variables are set
if [ -z "$DB_USERNAME" ] || [ -z "$DB_DATABASE_NAME" ] || [ -z "$UPLOAD_LOCATION" ]; then
    echo "Error: Required environment variables (DB_USERNAME, DB_DATABASE_NAME, UPLOAD_LOCATION) not found in .env file"
    echo "Please ensure your .env file contains:"
    echo "  - DB_USERNAME=<database_username>"
    echo "  - DB_DATABASE_NAME=<database_name>"
    echo "  - UPLOAD_LOCATION=<path_to_upload_directory>"
    exit 1
fi

# Notification function (matches backup script pattern)
send_notification() {
    local title="$1"
    local message="$2"
    local status="${3:-info}"  # success, error, warning, info
    local hostname=$(hostname)

    # Console notification
    log_message "$title: $message"

    # Webhook notification
    if [ -n "$WEBHOOK_URL" ]; then
        local tags="restore,immich,${hostname}"
        [ "$status" == "error" ] && tags="${tags},errors"
        [ "$status" == "warning" ] && tags="${tags},warnings"

        # Clean message without newlines or timestamps for webhook
        local webhook_message="$message"

        curl -s \
            -H "tags:${tags}" \
            -d "$webhook_message" \
            "$WEBHOOK_URL" 2>/dev/null || log_message "Warning: Failed to send webhook notification"
    fi
}

# Cleanup function to ensure containers are restarted
cleanup() {
    local exit_code=$?
    log_message "Running cleanup..."

    # Restart containers if they were stopped
    if [ "$CONTAINERS_STOPPED" = true ]; then
        log_message "Restarting Immich containers..."
        docker start immich_postgres 2>/dev/null || true
        docker start immich_server 2>/dev/null || true
        sleep 5

        # Check if containers started successfully
        if docker ps -q --filter "name=immich_postgres" | grep -q . && \
           docker ps -q --filter "name=immich_server" | grep -q .; then
            log_message "✅ Immich containers restarted successfully"
        else
            log_message "⚠️ Warning: Some containers may not have restarted properly"
        fi
    fi

    if [ $exit_code -ne 0 ]; then
        log_message "Restore failed with exit code $exit_code"
        send_notification "🚨 Immich Restore Failed" "Restoration process encountered an error (exit code: $exit_code)" "error"
    fi

    exit $exit_code
}

# Set up trap to call cleanup function on script exit
trap cleanup EXIT SIGINT SIGTERM

# Validate backup files integrity
validate_backup_files() {
    log_message "Validating backup files..."

    if [ "$SKIP_DB" = false ]; then
        if [ ! -f "$DB_BACKUP" ]; then
            log_message "Error: Database backup file not found: $DB_BACKUP"
            exit 1
        fi

        # Check if it's a gzipped file
        if [[ "$DB_BACKUP" == *.gz ]]; then
            if ! gunzip -t "$DB_BACKUP" 2>/dev/null; then
                log_message "Error: Database backup file appears to be corrupted (gzip test failed)"
                exit 1
            fi
        fi

        # Check available disk space for database restoration
        local db_size_bytes=$(stat -c%s "$DB_BACKUP" 2>/dev/null || echo "0")
        local available_space=$(df /tmp | tail -1 | awk '{print $4}')
        available_space=$((available_space * 1024))  # Convert to bytes

        # Estimate uncompressed size (conservative: 5x compressed size)
        local estimated_uncompressed=$((db_size_bytes * 5))

        if [ $estimated_uncompressed -gt $available_space ]; then
            log_message "Error: Insufficient disk space for database restoration"
            log_message "Estimated space needed: $(( estimated_uncompressed / 1024 / 1024 ))MB, Available: $(( available_space / 1024 / 1024 ))MB"
            exit 1
        fi

        local db_size=$(du -sh "$DB_BACKUP" 2>/dev/null | cut -f1 || echo "unknown")
        log_message "✓ Database backup validated: $DB_BACKUP ($db_size)"
    fi

    if [ "$SKIP_UPLOADS" = false ]; then
        if [ ! -f "$UPLOADS_BACKUP" ]; then
            log_message "Error: Uploads backup file not found: $UPLOADS_BACKUP"
            exit 1
        fi

        # Basic tar file validation
        if ! tar -tzf "$UPLOADS_BACKUP" >/dev/null 2>&1; then
            log_message "Error: Uploads backup file appears to be corrupted (tar test failed)"
            exit 1
        fi

        # Check available disk space for uploads restoration
        local uploads_size_bytes=$(stat -c%s "$UPLOADS_BACKUP" 2>/dev/null || echo "0")
        local upload_dir_space=$(df "$(dirname "$UPLOAD_LOCATION")" | tail -1 | awk '{print $4}')
        upload_dir_space=$((upload_dir_space * 1024))  # Convert to bytes

        # Estimate uncompressed size (conservative: 2x compressed size for media files)
        local estimated_uploads_uncompressed=$((uploads_size_bytes * 2))

        if [ $estimated_uploads_uncompressed -gt $upload_dir_space ]; then
            log_message "Error: Insufficient disk space for uploads restoration"
            log_message "Estimated space needed: $(( estimated_uploads_uncompressed / 1024 / 1024 ))MB, Available: $(( upload_dir_space / 1024 / 1024 ))MB"
            exit 1
        fi

        local uploads_size=$(du -sh "$UPLOADS_BACKUP" 2>/dev/null | cut -f1 || echo "unknown")
        log_message "✓ Uploads backup validated: $UPLOADS_BACKUP ($uploads_size)"
    fi
}

# Check container status and stop if needed
manage_containers() {
    local action="$1"  # stop or start

    if [ "$action" = "stop" ]; then
        log_message "Stopping Immich containers for restoration..."

        # Check and stop immich_server
        if docker ps -q --filter "name=immich_server" | grep -q .; then
            log_message "Stopping immich_server container..."
            docker stop immich_server || {
                log_message "Warning: Failed to stop immich_server container"
            }
        else
            log_message "immich_server container not running"
        fi

        # Check and stop immich_postgres
        if docker ps -q --filter "name=immich_postgres" | grep -q .; then
            log_message "Stopping immich_postgres container..."
            docker stop immich_postgres || {
                log_message "Error: Failed to stop immich_postgres container"
                exit 1
            }
        else
            log_message "immich_postgres container not running"
        fi

        # Wait for containers to fully stop
        sleep 5
        CONTAINERS_STOPPED=true

    elif [ "$action" = "start" ]; then
        log_message "Starting Immich containers..."

        # Start postgres first
        docker start immich_postgres || {
            log_message "Error: Failed to start immich_postgres container"
            exit 1
        }

        # Wait for postgres to be ready
        log_message "Waiting for PostgreSQL to be ready..."
        local max_attempts=15
        local attempt=0

        while [ $attempt -lt $max_attempts ]; do
            if docker exec immich_postgres pg_isready -U "$DB_USERNAME" >/dev/null 2>&1; then
                log_message "✓ PostgreSQL is ready"
                break
            fi
            attempt=$((attempt + 1))
            sleep 2
        done

        if [ $attempt -eq $max_attempts ]; then
            log_message "Warning: PostgreSQL may not be fully ready"
        fi

        # Start immich_server
        docker start immich_server || {
            log_message "Warning: Failed to start immich_server container"
        }

        # Wait for services to be ready
        sleep 5
        CONTAINERS_STOPPED=false
    fi
}

# Restore database
restore_database() {
    if [ "$SKIP_DB" = true ]; then
        log_message "Skipping database restoration (--skip-db)"
        return 0
    fi

    log_message "=== RESTORING DATABASE ==="

    # Create temporary directory for decompression if needed
    local temp_dir=$(mktemp -d)
    local sql_file="$DB_BACKUP"

    # Decompress if it's a gzipped file
    if [[ "$DB_BACKUP" == *.gz ]]; then
        log_message "Decompressing database backup..."
        sql_file="${temp_dir}/$(basename "${DB_BACKUP%.gz}")"

        if ! gunzip -c "$DB_BACKUP" > "$sql_file"; then
            log_message "Error: Failed to decompress database backup"
            rm -rf "$temp_dir"
            exit 1
        fi

        log_message "✓ Database backup decompressed to: $sql_file"
    fi

    # Start postgres container temporarily for restoration
    log_message "Starting PostgreSQL container for restoration..."
    docker start immich_postgres || {
        log_message "Error: Failed to start immich_postgres container"
        rm -rf "$temp_dir"
        exit 1
    }

    # Wait for PostgreSQL to be ready
    log_message "Waiting for PostgreSQL to be ready..."
    local max_attempts=30
    local attempt=0

    while [ $attempt -lt $max_attempts ]; do
        if docker exec immich_postgres pg_isready -U "$DB_USERNAME" >/dev/null 2>&1; then
            log_message "✓ PostgreSQL is ready (attempt $((attempt + 1)))"
            break
        fi
        attempt=$((attempt + 1))
        log_message "Waiting for PostgreSQL... (attempt $attempt/$max_attempts)"
        sleep 2
    done

    if [ $attempt -eq $max_attempts ]; then
        log_message "Error: PostgreSQL did not become ready within timeout (${max_attempts} attempts)"
        rm -rf "$temp_dir"
        exit 1
    fi

    # Drop existing database and recreate (if it exists)
    log_message "Preparing database for restoration..."

    # Validate database and username names to prevent injection
    if [[ ! "$DB_DATABASE_NAME" =~ ^[a-zA-Z0-9_]+$ ]] || [[ ! "$DB_USERNAME" =~ ^[a-zA-Z0-9_]+$ ]]; then
        log_message "Error: Invalid characters in database name or username"
        rm -rf "$temp_dir"
        exit 1
    fi

    docker exec immich_postgres psql -U "$DB_USERNAME" -c "DROP DATABASE IF EXISTS \"$DB_DATABASE_NAME\";" 2>/dev/null || true
    docker exec immich_postgres psql -U "$DB_USERNAME" -c "CREATE DATABASE \"$DB_DATABASE_NAME\";" || {
        log_message "Error: Failed to create database $DB_DATABASE_NAME"
        rm -rf "$temp_dir"
        exit 1
    }

    # Restore database
    log_message "Restoring database from backup..."
    log_message "This may take several minutes depending on database size..."

    if docker exec -i immich_postgres psql -U "$DB_USERNAME" -d "$DB_DATABASE_NAME" < "$sql_file"; then
        log_message "✅ Database restoration completed successfully"
    else
        log_message "Error: Database restoration failed"
        rm -rf "$temp_dir"
        exit 1
    fi

    # Stop postgres container
    docker stop immich_postgres

    # Cleanup temporary files
    rm -rf "$temp_dir"
}

# Restore uploads directory
restore_uploads() {
    if [ "$SKIP_UPLOADS" = true ]; then
        log_message "Skipping uploads restoration (--skip-uploads)"
        return 0
    fi

    log_message "=== RESTORING UPLOADS DIRECTORY ==="

    # Verify upload location exists or create it
    if [ ! -d "$UPLOAD_LOCATION" ]; then
        log_message "Creating upload directory: $UPLOAD_LOCATION"
        if ! mkdir -p "$UPLOAD_LOCATION"; then
            log_message "Error: Failed to create upload directory: $UPLOAD_LOCATION"
            exit 1
        fi
    fi

    # Backup existing content if any
    if [ "$(ls -A "$UPLOAD_LOCATION" 2>/dev/null)" ]; then
        local backup_existing="${UPLOAD_LOCATION}_backup_$(date +%Y%m%d_%H%M%S)"
        log_message "Backing up existing uploads to: $backup_existing"

        if ! mv "$UPLOAD_LOCATION" "$backup_existing"; then
            log_message "Error: Failed to backup existing uploads directory"
            exit 1
        fi

        # Recreate the upload directory
        mkdir -p "$UPLOAD_LOCATION"
    fi

    # Extract uploads backup with security safeguards
    log_message "Extracting uploads backup..."
    log_message "This may take a while depending on the size of your media library..."

    # Extract to parent directory with security options
    local parent_dir=$(dirname "$UPLOAD_LOCATION")
    local upload_dirname=$(basename "$UPLOAD_LOCATION")

    # Create a secure temporary extraction directory
    local temp_extract_dir=$(mktemp -d)

    # Extract with safety options to prevent path traversal
    if tar --no-absolute-names --strip-components=0 -xzf "$UPLOADS_BACKUP" -C "$temp_extract_dir"; then
        # Move extracted content to final location
        if [ -d "$UPLOAD_LOCATION" ]; then
            log_message "Backup existing upload directory..."
            mv "$UPLOAD_LOCATION" "${UPLOAD_LOCATION}.backup.$(date +%s)" || {
                log_message "Error: Failed to backup existing upload directory"
                rm -rf "$temp_extract_dir"
                exit 1
            }
        fi

        # Move extracted content to final location
        mv "$temp_extract_dir"/* "$UPLOAD_LOCATION" 2>/dev/null || {
            # Handle case where extraction created nested directories
            local extracted_dir=$(find "$temp_extract_dir" -mindepth 1 -maxdepth 1 -type d | head -1)
            if [ -n "$extracted_dir" ]; then
                mv "$extracted_dir" "$UPLOAD_LOCATION" || {
                    log_message "Error: Failed to move extracted files to upload location"
                    rm -rf "$temp_extract_dir"
                    exit 1
                }
            else
                log_message "Error: No content found in uploads backup"
                rm -rf "$temp_extract_dir"
                exit 1
            fi
        }

        rm -rf "$temp_extract_dir"
        log_message "✅ Uploads restoration completed successfully"
    else
        log_message "Error: Failed to extract uploads backup"
        rm -rf "$temp_extract_dir"
        exit 1
    fi

    # Set proper ownership and permissions
    log_message "Setting proper ownership and permissions..."

    # Find the user that should own these files (check docker container user)
    local container_user=$(docker exec immich_server id -u 2>/dev/null || echo "999")
    local container_group=$(docker exec immich_server id -g 2>/dev/null || echo "999")

    # Set ownership (use chown with numeric IDs to avoid user name conflicts)
    if command -v chown >/dev/null; then
        chown -R "$container_user:$container_group" "$UPLOAD_LOCATION" 2>/dev/null || {
            log_message "Warning: Could not set ownership, you may need to run: sudo chown -R $container_user:$container_group $UPLOAD_LOCATION"
        }
    fi

    # Set permissions
    find "$UPLOAD_LOCATION" -type d -exec chmod 755 {} \; 2>/dev/null || true
    find "$UPLOAD_LOCATION" -type f -exec chmod 644 {} \; 2>/dev/null || true

    log_message "✓ Ownership and permissions set"
}

# Verify restoration
verify_restoration() {
    log_message "=== VERIFYING RESTORATION ==="

    # Start containers for verification
    manage_containers start

    # Wait for services to be fully ready
    log_message "Waiting for Immich services to be ready..."
    sleep 15

    # Check database connectivity
    if [ "$SKIP_DB" = false ]; then
        log_message "Verifying database connectivity..."
        if docker exec immich_postgres psql -U "$DB_USERNAME" -d "$DB_DATABASE_NAME" -c "SELECT COUNT(*) FROM information_schema.tables;" >/dev/null 2>&1; then
            log_message "✅ Database connectivity verified"
        else
            log_message "⚠️ Warning: Database connectivity check failed"
        fi
    fi

    # Check uploads directory
    if [ "$SKIP_UPLOADS" = false ]; then
        log_message "Verifying uploads directory..."
        if [ -d "$UPLOAD_LOCATION" ] && [ "$(ls -A "$UPLOAD_LOCATION" 2>/dev/null)" ]; then
            local upload_size=$(du -sh "$UPLOAD_LOCATION" 2>/dev/null | cut -f1 || echo "unknown")
            log_message "✅ Uploads directory verified: $UPLOAD_LOCATION ($upload_size)"
        else
            log_message "⚠️ Warning: Uploads directory appears empty or inaccessible"
        fi
    fi

    # Check container status
    log_message "Verifying container status..."
    local containers_ok=true

    if ! docker ps -q --filter "name=immich_postgres" | grep -q .; then
        log_message "⚠️ Warning: immich_postgres container is not running"
        containers_ok=false
    fi

    if ! docker ps -q --filter "name=immich_server" | grep -q .; then
        log_message "⚠️ Warning: immich_server container is not running"
        containers_ok=false
    fi

    if [ "$containers_ok" = true ]; then
        log_message "✅ All containers are running"
    fi
}

# Initialize variables
CONTAINERS_STOPPED=false
RESTORE_START_TIME=$(date +%s)

# Main restoration logic
log_message "=== IMMICH RESTORATION STARTED ==="
send_notification "🔄 Immich Restore Started" "Beginning restoration from backup files" "info"

# Validate backup files first (even in dry-run mode for safety)
validate_backup_files

if [ "$DRY_RUN" = true ]; then
    log_message "=== DRY RUN MODE - NO ACTUAL RESTORATION WILL BE PERFORMED ==="
    log_message ""
    log_message "Configuration:"
    log_message "  - Database: ${DB_DATABASE_NAME}"
    log_message "  - Username: ${DB_USERNAME}"
    log_message "  - Upload Location: ${UPLOAD_LOCATION}"
    log_message "  - Database backup: ${DB_BACKUP:-SKIPPED}"
    log_message "  - Uploads backup: ${UPLOADS_BACKUP:-SKIPPED}"
    log_message ""
    log_message "Restoration process would:"
    [ "$SKIP_DB" = false ] && log_message "  1. Stop Immich containers"
    [ "$SKIP_DB" = false ] && log_message "  2. Restore database from: $DB_BACKUP"
    [ "$SKIP_UPLOADS" = false ] && log_message "  3. Restore uploads to: $UPLOAD_LOCATION"
    log_message "  4. Restart containers and verify"
    log_message ""
    log_message "=== DRY RUN COMPLETE - No changes would be made ==="
    exit 0
fi

# Stop containers
manage_containers stop

# Perform restoration
restore_database
restore_uploads

# Verify restoration
verify_restoration

# Calculate restoration time and send success notification
RESTORE_END_TIME=$(date +%s)
TOTAL_RESTORE_TIME=$((RESTORE_END_TIME - RESTORE_START_TIME))
RESTORE_TIME_FORMATTED=""

if [ $TOTAL_RESTORE_TIME -gt 3600 ]; then
    HOURS=$((TOTAL_RESTORE_TIME / 3600))
    MINUTES=$(((TOTAL_RESTORE_TIME % 3600) / 60))
    RESTORE_TIME_FORMATTED="${HOURS}h ${MINUTES}m"
elif [ $TOTAL_RESTORE_TIME -gt 60 ]; then
    MINUTES=$((TOTAL_RESTORE_TIME / 60))
    SECONDS=$((TOTAL_RESTORE_TIME % 60))
    RESTORE_TIME_FORMATTED="${MINUTES}m ${SECONDS}s"
else
    RESTORE_TIME_FORMATTED="${TOTAL_RESTORE_TIME}s"
fi

# Format file information for notification
NOTIFICATION_MESSAGE=""
if [ "$SKIP_DB" = false ]; then
    DB_FILENAME=$(basename "$DB_BACKUP")
    DB_SIZE=$(du -sh "$DB_BACKUP" 2>/dev/null | cut -f1 || echo "unknown")
    NOTIFICATION_MESSAGE="📦 Database: ${DB_FILENAME} (${DB_SIZE})"
fi

if [ "$SKIP_UPLOADS" = false ]; then
    UPLOADS_FILENAME=$(basename "$UPLOADS_BACKUP")
    UPLOADS_SIZE=$(du -sh "$UPLOADS_BACKUP" 2>/dev/null | cut -f1 || echo "unknown")
    if [ -n "$NOTIFICATION_MESSAGE" ]; then
        NOTIFICATION_MESSAGE="${NOTIFICATION_MESSAGE}\n📁 Uploads: ${UPLOADS_FILENAME} (${UPLOADS_SIZE})"
    else
        NOTIFICATION_MESSAGE="📁 Uploads: ${UPLOADS_FILENAME} (${UPLOADS_SIZE})"
    fi
fi

NOTIFICATION_MESSAGE="${NOTIFICATION_MESSAGE}\n✅ Restoration completed successfully in ${RESTORE_TIME_FORMATTED}"
NOTIFICATION_MESSAGE="${NOTIFICATION_MESSAGE}\n🏠 Restored to: $(hostname)"

send_notification "✅ Immich Restore Completed" "$NOTIFICATION_MESSAGE" "success"

log_message "=== IMMICH RESTORATION COMPLETED SUCCESSFULLY ==="
log_message "Database restored to: $DB_DATABASE_NAME"
log_message "Uploads restored to: $UPLOAD_LOCATION"
log_message "All containers are running and verified"