#!/bin/bash
###############################################################################
# mockupAWS Database Restore Script v1.0.0
#
# Description: PostgreSQL database restore with Point-in-Time Recovery support
#
# Features:
#   - Full database restore from backup
#   - Point-in-Time Recovery (PITR)
#   - Integrity verification
#   - Decryption support
#   - S3 download
#
# Recovery Objectives:
#   - RTO (Recovery Time Objective): < 1 hour
#   - RPO (Recovery Point Objective): < 5 minutes
#
# Usage:
#   ./scripts/restore.sh latest                    # Restore latest backup
#   ./scripts/restore.sh s3://bucket/key           # Restore from S3
#   ./scripts/restore.sh /path/to/backup.enc       # Restore from local file
#   ./scripts/restore.sh latest --target-time "2026-04-07 14:30:00"  # PITR
#   ./scripts/restore.sh latest --dry-run          # Verify without restoring
#
# Environment Variables:
#   DATABASE_URL           - Target PostgreSQL connection (required)
#   BACKUP_ENCRYPTION_KEY  - AES-256 decryption key
#   BACKUP_BUCKET          - S3 bucket name
#   AWS_ACCESS_KEY_ID      - AWS credentials
#   AWS_SECRET_ACCESS_KEY  - AWS credentials
#
###############################################################################

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
RESTORE_DIR="${PROJECT_ROOT}/storage/restore"
LOG_DIR="${PROJECT_ROOT}/storage/logs"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

# Default values
TARGET_TIME=""
DRY_RUN=false
VERIFY_ONLY=false
SKIP_BACKUP=false

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging
log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] ✓${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] ⚠${NC} $1"
}

log_error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ✗${NC} $1"
}

# Create directories
mkdir -p "$RESTORE_DIR" "$LOG_DIR"

# Validate environment
validate_env() {
    local missing=()
    
    if [[ -z "${DATABASE_URL:-}" ]]; then
        missing+=("DATABASE_URL")
    fi
    
    if [[ ${#missing[@]} -gt 0 ]]; then
        log_error "Missing required environment variables: ${missing[*]}"
        exit 1
    fi
    
    if [[ -z "${BACKUP_ENCRYPTION_KEY:-}" ]]; then
        log_warn "BACKUP_ENCRYPTION_KEY not set - assuming unencrypted backups"
    fi
}

# Parse database URL
parse_database_url() {
    local url="$1"
    
    # Remove protocol
    local conn="${url#postgresql://}"
    conn="${conn#postgresql+asyncpg://}"
    conn="${conn#postgres://}"
    
    # Parse user:password@host:port/database
    if [[ "$conn" =~ ^([^:]+):([^@]+)@([^:]+):?([0-9]*)/([^?]+) ]]; then
        DB_USER="${BASH_REMATCH[1]}"
        DB_PASS="${BASH_REMATCH[2]}"
        DB_HOST="${BASH_REMATCH[3]}"
        DB_PORT="${BASH_REMATCH[4]:-5432}"
        DB_NAME="${BASH_REMATCH[5]}"
    else
        log_error "Could not parse DATABASE_URL"
        exit 1
    fi
    
    export PGPASSWORD="$DB_PASS"
}

# Decrypt file
decrypt_file() {
    local input_file="$1"
    local output_file="$2"
    
    if [[ -n "${BACKUP_ENCRYPTION_KEY:-}" ]]; then
        log "Decrypting backup..."
        openssl enc -aes-256-cbc -d -pbkdf2 \
            -in "$input_file" \
            -out "$output_file" \
            -pass pass:"$BACKUP_ENCRYPTION_KEY" 2>/dev/null || {
            log_error "Decryption failed - check encryption key"
            exit 1
        }
        log_success "Decryption completed"
    else
        cp "$input_file" "$output_file"
    fi
}

# Download from S3
download_from_s3() {
    local s3_url="$1"
    local output_file="$2"
    
    log "Downloading from S3: $s3_url"
    aws s3 cp "$s3_url" "$output_file" || {
        log_error "Failed to download from S3"
        exit 1
    }
    log_success "Download completed"
}

# Find latest backup
find_latest_backup() {
    local backup_bucket="${BACKUP_BUCKET:-}"
    
    if [[ -z "$backup_bucket" ]]; then
        # Look for local backups
        local latest_backup
        latest_backup=$(ls -t "$RESTORE_DIR"/../backups/mockupaws_full_*.sql.gz.enc 2>/dev/null | head -1)
        
        if [[ -z "$latest_backup" ]]; then
            log_error "No local backups found"
            exit 1
        fi
        
        echo "$latest_backup"
    else
        # Find latest in S3
        local latest_key
        latest_key=$(aws s3 ls "s3://$backup_bucket/backups/full/" --recursive | \
            grep "mockupaws_full_.*\.sql\.gz\.enc$" | \
            sort | tail -1 | awk '{print $4}')
        
        if [[ -z "$latest_key" ]]; then
            log_error "No backups found in S3"
            exit 1
        fi
        
        echo "s3://$backup_bucket/$latest_key"
    fi
}

# Verify backup integrity
verify_backup() {
    local backup_file="$1"
    
    log "Verifying backup integrity..."
    
    # Decrypt to temp file
    local temp_decrypted="${RESTORE_DIR}/verify_${TIMESTAMP}.tmp"
    decrypt_file "$backup_file" "$temp_decrypted"
    
    # Decompress
    local temp_sql="${RESTORE_DIR}/verify_${TIMESTAMP}.sql"
    gunzip -c "$temp_decrypted" > "$temp_sql" 2>/dev/null || {
        # Might not be compressed
        mv "$temp_decrypted" "$temp_sql"
    }
    
    # Verify with pg_restore
    if pg_restore --list "$temp_sql" > /dev/null 2>&1; then
        local object_count
        object_count=$(pg_restore --list "$temp_sql" | wc -l)
        log_success "Backup verification passed"
        log "  Objects in backup: $object_count"
        rm -f "$temp_sql" "$temp_decrypted"
        return 0
    else
        log_error "Backup verification failed - file may be corrupted"
        rm -f "$temp_sql" "$temp_decrypted"
        return 1
    fi
}

# Pre-restore checks
pre_restore_checks() {
    log "Performing pre-restore checks..."
    
    # Check if target database exists
    if psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="postgres" \
        --command="SELECT 1 FROM pg_database WHERE datname = '$DB_NAME';" \
        --tuples-only --no-align 2>/dev/null | grep -q 1; then
        
        log_warn "Target database '$DB_NAME' exists"
        
        if [[ "$SKIP_BACKUP" == false ]]; then
            log "Creating safety backup of existing database..."
            local safety_backup="${RESTORE_DIR}/safety_backup_${TIMESTAMP}.sql"
            pg_dump \
                --host="$DB_HOST" \
                --port="$DB_PORT" \
                --username="$DB_USER" \
                --dbname="$DB_NAME" \
                --format=plain \
                --file="$safety_backup" \
                2>/dev/null || log_warn "Could not create safety backup"
        fi
    fi
    
    # Check disk space
    local available_space
    available_space=$(df -k "$RESTORE_DIR" | awk 'NR==2 {print $4}')
    local required_space=1048576  # 1GB in KB
    
    if [[ $available_space -lt $required_space ]]; then
        log_error "Insufficient disk space (need ~1GB, have ${available_space}KB)"
        exit 1
    fi
    
    log_success "Pre-restore checks passed"
}

# Restore database
restore_database() {
    local backup_file="$1"
    
    log "Starting database restore..."
    
    if [[ "$DRY_RUN" == true ]]; then
        log_warn "DRY RUN MODE - No actual changes will be made"
        verify_backup "$backup_file"
        log_success "Dry run completed successfully"
        return 0
    fi
    
    # Verify first
    if ! verify_backup "$backup_file"; then
        log_error "Backup verification failed - aborting restore"
        exit 1
    fi
    
    # Decrypt
    local decrypted_file="${RESTORE_DIR}/restore_${TIMESTAMP}.sql.gz"
    decrypt_file "$backup_file" "$decrypted_file"
    
    # Drop and recreate database
    log "Dropping existing database (if exists)..."
    psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="postgres" \
        --command="DROP DATABASE IF EXISTS \"$DB_NAME\";" \
        2>/dev/null || true
    
    log "Creating new database..."
    psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="postgres" \
        --command="CREATE DATABASE \"$DB_NAME\";" \
        2>/dev/null
    
    # Restore
    log "Restoring database from backup..."
    pg_restore \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="$DB_NAME" \
        --jobs=4 \
        --verbose \
        "$decrypted_file" \
        2>"${LOG_DIR}/restore_${TIMESTAMP}.log" || {
        log_warn "pg_restore completed with warnings (check log)"
    }
    
    # Cleanup
    rm -f "$decrypted_file"
    
    log_success "Database restore completed"
}

# Point-in-Time Recovery
restore_pitr() {
    local backup_file="$1"
    local target_time="$2"
    
    log "Starting Point-in-Time Recovery to: $target_time"
    log_warn "PITR requires WAL archiving to be configured"
    
    if [[ "$DRY_RUN" == true ]]; then
        log "Would recover to: $target_time"
        return 0
    fi
    
    # This is a simplified PITR - in production, use proper WAL archiving
    restore_database "$backup_file"
    
    # Apply WAL files up to target time
    log "Applying WAL files up to $target_time..."
    
    # Note: Full PITR implementation requires:
    # 1. archive_command configured in PostgreSQL
    # 2. restore_command configured
    # 3. recovery_target_time set
    # 4. Recovery mode trigger file
    
    log_warn "PITR implementation requires manual WAL replay configuration"
    log "Refer to docs/BACKUP-RESTORE.md for detailed PITR procedures"
}

# Post-restore validation
post_restore_validation() {
    log "Performing post-restore validation..."
    
    # Check database is accessible
    local table_count
    table_count=$(psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="$DB_NAME" \
        --command="SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';" \
        --tuples-only --no-align 2>/dev/null)
    
    if [[ -z "$table_count" ]] || [[ "$table_count" == "0" ]]; then
        log_error "Post-restore validation failed - no tables found"
        exit 1
    fi
    
    log "  Tables restored: $table_count"
    
    # Check key tables
    local key_tables=("scenarios" "scenario_logs" "scenario_metrics" "users" "reports")
    for table in "${key_tables[@]}"; do
        if psql \
            --host="$DB_HOST" \
            --port="$DB_PORT" \
            --username="$DB_USER" \
            --dbname="$DB_NAME" \
            --command="SELECT 1 FROM $table LIMIT 1;" \
            > /dev/null 2>&1; then
            log_success "  Table '$table' accessible"
        else
            log_warn "  Table '$table' not accessible or empty"
        fi
    done
    
    # Record restore in database
    psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USER" \
        --dbname="$DB_NAME" \
        --command="
            CREATE TABLE IF NOT EXISTS restore_history (
                id SERIAL PRIMARY KEY,
                restored_at TIMESTAMP DEFAULT NOW(),
                source_backup TEXT,
                target_time TIMESTAMP,
                table_count INTEGER,
                status VARCHAR(50)
            );
            INSERT INTO restore_history (source_backup, target_time, table_count, status)
            VALUES ('$BACKUP_SOURCE', '$TARGET_TIME', $table_count, 'completed');
        " \
        2>/dev/null || true
    
    log_success "Post-restore validation completed"
}

# Print restore summary
print_summary() {
    local start_time="$1"
    local end_time
    end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    echo ""
    echo "=============================================="
    echo "  RESTORE SUMMARY"
    echo "=============================================="
    echo "  Source: $BACKUP_SOURCE"
    echo "  Target: $DATABASE_URL"
    echo "  Duration: ${duration}s"
    if [[ -n "$TARGET_TIME" ]]; then
        echo "  PITR Target: $TARGET_TIME"
    fi
    echo "  Log file: ${LOG_DIR}/restore_${TIMESTAMP}.log"
    echo "=============================================="
}

# Main restore function
main() {
    local backup_source="$1"
    shift
    
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --target-time)
                TARGET_TIME="$2"
                shift 2
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            --verify-only)
                VERIFY_ONLY=true
                shift
                ;;
            --skip-backup)
                SKIP_BACKUP=true
                shift
                ;;
            *)
                shift
                ;;
        esac
    done
    
    local start_time
    start_time=$(date +%s)
    BACKUP_SOURCE="$backup_source"
    
    validate_env
    parse_database_url "$DATABASE_URL"
    
    log "mockupAWS Database Restore v1.0.0"
    log "=================================="
    
    # Resolve backup source
    local backup_file
    if [[ "$backup_source" == "latest" ]]; then
        backup_file=$(find_latest_backup)
        log "Latest backup: $backup_file"
    elif [[ "$backup_source" == s3://* ]]; then
        backup_file="${RESTORE_DIR}/download_${TIMESTAMP}.sql.gz.enc"
        download_from_s3 "$backup_source" "$backup_file"
    elif [[ -f "$backup_source" ]]; then
        backup_file="$backup_source"
    else
        log_error "Invalid backup source: $backup_source"
        exit 1
    fi
    
    if [[ "$VERIFY_ONLY" == true ]]; then
        verify_backup "$backup_file"
        exit 0
    fi
    
    pre_restore_checks
    
    if [[ -n "$TARGET_TIME" ]]; then
        restore_pitr "$backup_file" "$TARGET_TIME"
    else
        restore_database "$backup_file"
    fi
    
    post_restore_validation
    
    print_summary "$start_time"
    
    log_success "Restore completed successfully!"
    
    # Cleanup downloaded S3 files
    if [[ "$backup_file" == "${RESTORE_DIR}/download_"* ]]; then
        rm -f "$backup_file"
    fi
}

# Show usage
usage() {
    echo "mockupAWS Database Restore Script v1.0.0"
    echo ""
    echo "Usage: $0 <backup-source> [options]"
    echo ""
    echo "Backup Sources:"
    echo "  latest                    Restore latest backup from S3 or local"
    echo "  s3://bucket/path          Restore from S3 URL"
    echo "  /path/to/backup.enc       Restore from local file"
    echo ""
    echo "Options:"
    echo "  --target-time 'YYYY-MM-DD HH:MM:SS'  Point-in-Time Recovery"
    echo "  --dry-run                 Verify backup without restoring"
    echo "  --verify-only             Only verify backup integrity"
    echo "  --skip-backup             Skip safety backup of existing DB"
    echo ""
    echo "Environment Variables:"
    echo "  DATABASE_URL           - Target PostgreSQL connection (required)"
    echo "  BACKUP_ENCRYPTION_KEY  - AES-256 decryption key"
    echo "  BACKUP_BUCKET          - S3 bucket name"
    echo ""
    echo "Examples:"
    echo "  $0 latest"
    echo "  $0 latest --target-time '2026-04-07 14:30:00'"
    echo "  $0 s3://mybucket/backups/full/20260407/backup.enc"
    echo "  $0 /backups/mockupaws_full_20260407_120000.sql.gz.enc --dry-run"
    echo ""
}

# Main entry point
if [[ $# -eq 0 ]]; then
    usage
    exit 1
fi

main "$@"