#!/usr/bin/env bash
#
# LLMHosts.com Database Restore Script
# Restores SQLite database from backup with integrity verification
#
# Usage:
#   ./restore.sh <backup-file> [--verify-only]
#
# Examples:
#   ./restore.sh /backups/llmhost_20260218_120000.db.gz
#   ./restore.sh s3://backup-bucket/llmhost_20260218_120000.db.gz
#   ./restore.sh /backups/llmhost_20260218_120000.db.gz --verify-only

set -euo pipefail

# Configuration
DB_PATH="${DB_PATH:-/data/llmhost.db}"
TEMP_DIR=$(mktemp -d)
VERIFY_ONLY=false

# Parse arguments
if [[ $# -lt 1 ]]; then
  echo "Usage: $0 <backup-file> [--verify-only]"
  exit 1
fi

BACKUP_SOURCE="$1"
shift

while [[ $# -gt 0 ]]; do
  case $1 in
    --verify-only)
      VERIFY_ONLY=true
      shift
      ;;
    *)
      echo "Unknown option: $1"
      exit 1
      ;;
  esac
done

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

error() {
  echo "[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $*" >&2
}

# Cleanup on exit
cleanup() {
  rm -rf "$TEMP_DIR"
}
trap cleanup EXIT

# Download backup if remote
BACKUP_FILE="$TEMP_DIR/backup.db.gz"

if [[ "$BACKUP_SOURCE" =~ ^s3:// ]]; then
  log "Downloading from S3: $BACKUP_SOURCE"
  aws s3 cp "$BACKUP_SOURCE" "$BACKUP_FILE"
  CHECKSUM_SOURCE="${BACKUP_SOURCE%.gz}.sha256"
  aws s3 cp "$CHECKSUM_SOURCE" "$TEMP_DIR/backup.sha256" || true
elif [[ "$BACKUP_SOURCE" =~ ^b2:// ]]; then
  log "Downloading from B2: $BACKUP_SOURCE"
  BUCKET=$(echo "$BACKUP_SOURCE" | cut -d/ -f3)
  FILENAME=$(echo "$BACKUP_SOURCE" | cut -d/ -f4-)
  b2 download-file-by-name --quiet "$BUCKET" "$FILENAME" "$BACKUP_FILE"
  b2 download-file-by-name --quiet "$BUCKET" "${FILENAME%.gz}.sha256" "$TEMP_DIR/backup.sha256" || true
else
  log "Using local backup: $BACKUP_SOURCE"
  cp "$BACKUP_SOURCE" "$BACKUP_FILE"
  if [[ -f "${BACKUP_SOURCE%.gz}.sha256" ]]; then
    cp "${BACKUP_SOURCE%.gz}.sha256" "$TEMP_DIR/backup.sha256"
  fi
fi

# Verify checksum if available
if [[ -f "$TEMP_DIR/backup.sha256" ]]; then
  log "Verifying checksum"
  EXPECTED_CHECKSUM=$(cat "$TEMP_DIR/backup.sha256")
  ACTUAL_CHECKSUM=$(sha256sum "$BACKUP_FILE" | awk '{print $1}')

  if [[ "$EXPECTED_CHECKSUM" == "$ACTUAL_CHECKSUM" ]]; then
    log "Checksum verified: $ACTUAL_CHECKSUM"
  else
    error "Checksum mismatch!"
    error "Expected: $EXPECTED_CHECKSUM"
    error "Actual:   $ACTUAL_CHECKSUM"
    exit 1
  fi
else
  log "Warning: No checksum file found, skipping verification"
fi

# Decompress backup
log "Decompressing backup"
gunzip -c "$BACKUP_FILE" > "$TEMP_DIR/backup.db"

# Verify SQLite integrity
log "Verifying SQLite integrity"
if sqlite3 "$TEMP_DIR/backup.db" "PRAGMA integrity_check;" | grep -q "ok"; then
  log "Backup integrity verified"
else
  error "Backup integrity check failed!"
  exit 1
fi

# Verify schema version matches current
BACKUP_VERSION=$(sqlite3 "$TEMP_DIR/backup.db" "PRAGMA user_version;" 2>/dev/null || echo "0")
if [[ -f "$DB_PATH" ]]; then
  CURRENT_VERSION=$(sqlite3 "$DB_PATH" "PRAGMA user_version;" 2>/dev/null || echo "0")
  log "Backup schema version: $BACKUP_VERSION"
  log "Current schema version: $CURRENT_VERSION"

  if [[ "$BACKUP_VERSION" != "$CURRENT_VERSION" ]]; then
    log "Warning: Schema version mismatch (may need migration)"
  fi
fi

# If verify-only, exit here
if [[ "$VERIFY_ONLY" == "true" ]]; then
  log "Verification complete (verify-only mode, no restore performed)"
  exit 0
fi

# Backup current database (if exists)
if [[ -f "$DB_PATH" ]]; then
  CURRENT_BACKUP="${DB_PATH}.before-restore-$(date +%Y%m%d_%H%M%S)"
  log "Backing up current database to: $CURRENT_BACKUP"
  cp "$DB_PATH" "$CURRENT_BACKUP"
  log "Current database backed up (can rollback if needed)"
fi

# Restore database
log "Restoring database to: $DB_PATH"
cp "$TEMP_DIR/backup.db" "$DB_PATH"

# Verify restored database
log "Verifying restored database"
if sqlite3 "$DB_PATH" "PRAGMA integrity_check;" | grep -q "ok"; then
  log "Restore successful!"
else
  error "Restored database failed integrity check!"

  # Rollback if we backed up
  if [[ -n "${CURRENT_BACKUP:-}" ]]; then
    log "Rolling back to previous database"
    cp "$CURRENT_BACKUP" "$DB_PATH"
  fi

  exit 1
fi

# Summary
DB_SIZE=$(du -h "$DB_PATH" | cut -f1)
RECORD_COUNT=$(sqlite3 "$DB_PATH" "SELECT COUNT(*) FROM sqlite_master WHERE type='table';" || echo "unknown")
log "Restore complete: $DB_PATH (${DB_SIZE}, ${RECORD_COUNT} tables)"

# Prometheus metrics (if enabled)
if [[ -n "${METRICS_DIR:-}" ]]; then
  mkdir -p "$METRICS_DIR"
  echo "llmhost_restore_last_success_timestamp $(date +%s)" > "$METRICS_DIR/restore.prom"
fi

exit 0
