#!/usr/bin/env bash
set -euo pipefail

if [ $# -lt 1 ]; then
    echo "Usage: $0 <backup_filename>"
    echo "Example: $0 raseel_20260928_120000.dump.gz.enc"
    echo ""
    echo "Downloads the backup from S3, decrypts, decompresses, and restores."
    exit 1
fi

BACKUP_NAME="$1"
WORK_DIR="${WORK_DIR:-/tmp/raseel-restore}"
LOG_FILE="${LOG_FILE:-storage/logs/restore.log}"
AWS_BUCKET="${AWS_BUCKET:-raseel-backups}"
AWS_DEFAULT_REGION="${AWS_DEFAULT_REGION:-us-east-1}"

mkdir -p "$WORK_DIR"
mkdir -p "$(dirname "$LOG_FILE")"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

cleanup() {
    EXIT_CODE=$?
    rm -rf "$WORK_DIR" 2>/dev/null || true
    if [ $EXIT_CODE -ne 0 ]; then
        log "Restore failed with exit code ${EXIT_CODE}"
    fi
    exit $EXIT_CODE
}

trap cleanup EXIT

check_env() {
    local missing=0
    for var in DB_HOST DB_PORT DB_DATABASE DB_USERNAME DB_PASSWORD BACKUP_ENCRYPTION_PASSWORD AWS_BUCKET AWS_DEFAULT_REGION; do
        if [ -z "${!var:-}" ]; then
            log "ERROR: Missing environment variable: ${var}"
            missing=1
        fi
    done
    if [ -z "${AWS_ACCESS_KEY_ID:-}" ] || [ -z "${AWS_SECRET_ACCESS_KEY:-}" ]; then
        log "ERROR: Missing AWS credentials"
        missing=1
    fi
    return $missing
}

download_backup() {
    local s3_key="backups/$(date +%Y/%m)/${BACKUP_NAME}"
    local s3_uri="s3://${AWS_BUCKET}/${s3_key}"

    if ! aws s3 ls "$s3_uri" --region "$AWS_DEFAULT_REGION" >/dev/null 2>&1; then
        log "Backup not found at ${s3_uri}, searching..."
        local found_key
        found_key=$(aws s3 ls "s3://${AWS_BUCKET}/backups/" --recursive --region "$AWS_DEFAULT_REGION" \
            | grep "${BACKUP_NAME}" \
            | awk '{print $4}' \
            | head -1)
        if [ -z "$found_key" ]; then
            log "ERROR: Backup ${BACKUP_NAME} not found in S3"
            exit 1
        fi
        s3_uri="s3://${AWS_BUCKET}/${found_key}"
    fi

    log "Downloading ${s3_uri}"
    aws s3 cp "$s3_uri" "${WORK_DIR}/${BACKUP_NAME}" \
        --region "$AWS_DEFAULT_REGION" \
        --no-progress
}

decrypt_backup() {
    log "Decrypting backup"
    openssl enc -aes-256-cbc -pbkdf2 -d \
        -pass "env:BACKUP_ENCRYPTION_PASSWORD" \
        -in "${WORK_DIR}/${BACKUP_NAME}" \
        -out "${WORK_DIR}/${BACKUP_NAME}.gz"
}

decompress_backup() {
    log "Decompressing backup"
    gunzip -f "${WORK_DIR}/${BACKUP_NAME}.gz"
    local dump_file="${WORK_DIR}/${BACKUP_NAME%.gz}"
    mv "$dump_file" "${WORK_DIR}/restore.dump" 2>/dev/null || true
}

restore_database() {
    log "Restoring database ${DB_DATABASE}"

    PGPASSWORD="$DB_PASSWORD" pg_restore \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USERNAME" \
        --dbname="$DB_DATABASE" \
        --no-owner \
        --no-privileges \
        --clean \
        --if-exists \
        --verbose \
        "${WORK_DIR}/restore.dump" 2>&1 | tee -a "$LOG_FILE" || true
}

verify_restore() {
    log "Verifying restore"

    local table_count
    table_count=$(PGPASSWORD="$DB_PASSWORD" psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USERNAME" \
        --dbname="$DB_DATABASE" \
        --tuples-only \
        --no-align \
        --command="SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';")

    table_count=$(echo "$table_count" | tr -d '[:space:]')

    if [ -z "$table_count" ] || [ "$table_count" -eq 0 ]; then
        log "ERROR: Verification failed - no tables found in database"
        exit 1
    fi

    log "Verification passed: ${table_count} tables restored"

    local tenant_count
    tenant_count=$(PGPASSWORD="$DB_PASSWORD" psql \
        --host="$DB_HOST" \
        --port="$DB_PORT" \
        --username="$DB_USERNAME" \
        --dbname="$DB_DATABASE" \
        --tuples-only \
        --no-align \
        --command="SELECT COUNT(*) FROM tenants;" 2>/dev/null || echo "0")

    tenant_count=$(echo "$tenant_count" | tr -d '[:space:]')
    log "Tenant count in restored database: ${tenant_count}"
}

main() {
    log "=== Starting restore of ${BACKUP_NAME} ==="

    if ! check_env; then
        log "Environment validation failed"
        exit 1
    fi

    download_backup
    decrypt_backup
    decompress_backup
    restore_database
    verify_restore

    log "=== Restore completed successfully ==="
    exit 0
}

main "$@"