#!/usr/bin/env bash
#
# StemeDB Restore Script
#
# Restores WAL and database files from a backup created by backup-stemedb.sh.
#
# Usage:
#   ./scripts/restore-stemedb.sh backups/stemedb-backup-20260208-120000/
#   ./scripts/restore-stemedb.sh backups/stemedb-backup-*/ --force
#
# Safety:
#   - Checks that StemeDB is NOT running before restore
#   - Refuses to overwrite non-empty target dirs without --force
#   - With --force, renames existing dirs (never deletes)
#
# Exit codes:
#   0 - Restore completed successfully
#   1 - Restore failed
#

set -euo pipefail

# Configuration
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
readonly API_HOST="${STEMEDB_BIND_ADDR:-127.0.0.1:18180}"
readonly TIMESTAMP="$(date +%Y%m%d-%H%M%S)"

# Colors (if terminal supports it)
if [[ -t 1 ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[0;33m'
    BLUE='\033[0;34m'
    NC='\033[0m'
else
    RED=''
    GREEN=''
    YELLOW=''
    BLUE=''
    NC=''
fi

# Logging helpers
info() { echo -e "${BLUE}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC} $*"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $*"; }
fail() { echo -e "${RED}[FAIL]${NC} $*"; exit 1; }

# Defaults
BACKUP_PATH=""
TARGET_WAL="${STEMEDB_WAL_DIR:-${PROJECT_DIR}/data/wal}"
TARGET_DB="${STEMEDB_DB_DIR:-${PROJECT_DIR}/data/db}"
FORCE=false

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --target-wal)
            TARGET_WAL="$2"
            shift 2
            ;;
        --target-db)
            TARGET_DB="$2"
            shift 2
            ;;
        --force)
            FORCE=true
            shift
            ;;
        --help|-h)
            echo "Usage: $0 <backup-path> [--target-wal <dir>] [--target-db <dir>] [--force]"
            echo ""
            echo "Restore StemeDB from a backup."
            echo ""
            echo "Arguments:"
            echo "  <backup-path>      Path to backup directory (must contain backup-metadata.json)"
            echo ""
            echo "Options:"
            echo "  --target-wal <dir> Target WAL directory (default: data/wal)"
            echo "  --target-db <dir>  Target DB directory (default: data/db)"
            echo "  --force            Overwrite existing data (renames to .pre-restore-TIMESTAMP)"
            echo "  --help             Show this help message"
            echo ""
            echo "Environment:"
            echo "  STEMEDB_WAL_DIR    WAL directory (default: data/wal)"
            echo "  STEMEDB_DB_DIR     Database directory (default: data/db)"
            echo "  STEMEDB_BIND_ADDR  API address for running check (default: 127.0.0.1:18180)"
            exit 0
            ;;
        -*)
            fail "Unknown option: $1 (use --help for usage)"
            ;;
        *)
            if [[ -z "$BACKUP_PATH" ]]; then
                BACKUP_PATH="$1"
            else
                fail "Unexpected argument: $1"
            fi
            shift
            ;;
    esac
done

if [[ -z "$BACKUP_PATH" ]]; then
    fail "Backup path is required. Usage: $0 <backup-path> [--force]"
fi

main() {
    echo ""
    echo "=========================================="
    echo "  StemeDB Restore"
    echo "=========================================="
    echo ""

    # Validate backup
    if [[ ! -d "$BACKUP_PATH" ]]; then
        fail "Backup directory not found: ${BACKUP_PATH}"
    fi

    if [[ ! -f "${BACKUP_PATH}/backup-metadata.json" ]]; then
        fail "Not a valid backup: missing backup-metadata.json in ${BACKUP_PATH}"
    fi

    info "Backup: ${BACKUP_PATH}"
    info "Metadata:"
    cat "${BACKUP_PATH}/backup-metadata.json" | sed 's/^/  /'
    echo ""

    # Check StemeDB is NOT running
    info "Checking that StemeDB is not running..."
    if curl -s --connect-timeout 2 "http://${API_HOST}/v1/health" > /dev/null 2>&1; then
        fail "StemeDB is running at ${API_HOST}. Stop the server before restoring."
    fi
    success "StemeDB is not running"

    # Check what's in the backup
    local has_wal=false
    local has_db=false
    [[ -d "${BACKUP_PATH}/wal" ]] && has_wal=true
    [[ -d "${BACKUP_PATH}/db" ]] && has_db=true

    if [[ "$has_wal" == "false" ]]; then
        fail "Backup contains no WAL directory"
    fi

    # Handle existing target directories
    if [[ -d "$TARGET_WAL" && -n "$(ls -A "$TARGET_WAL" 2>/dev/null)" ]]; then
        if [[ "$FORCE" == "false" ]]; then
            fail "Target WAL directory is not empty: ${TARGET_WAL}\n  Use --force to rename existing data and proceed."
        fi
        local renamed="${TARGET_WAL}.pre-restore-${TIMESTAMP}"
        warn "Renaming existing WAL: ${TARGET_WAL} -> ${renamed}"
        mv "$TARGET_WAL" "$renamed"
    fi

    if [[ "$has_db" == "true" && -d "$TARGET_DB" && -n "$(ls -A "$TARGET_DB" 2>/dev/null)" ]]; then
        if [[ "$FORCE" == "false" ]]; then
            fail "Target DB directory is not empty: ${TARGET_DB}\n  Use --force to rename existing data and proceed."
        fi
        local renamed="${TARGET_DB}.pre-restore-${TIMESTAMP}"
        warn "Renaming existing DB: ${TARGET_DB} -> ${renamed}"
        mv "$TARGET_DB" "$renamed"
    fi

    # Restore WAL
    info "Restoring WAL..."
    mkdir -p "$TARGET_WAL"
    rsync -a "${BACKUP_PATH}/wal/" "${TARGET_WAL}/"
    local wal_files
    wal_files=$(find "$TARGET_WAL" -type f | wc -l)
    success "WAL restored: ${wal_files} files"

    # Verify WAL header magic bytes (STEM = first 4 bytes)
    local wal_valid=true
    for wal_file in "${TARGET_WAL}"/*.wal; do
        [[ -f "$wal_file" ]] || continue
        local magic
        magic=$(head -c 4 "$wal_file" | od -A n -t x1 | tr -d ' ')
        if [[ "$magic" == "5354454d" ]]; then
            success "WAL magic OK: $(basename "$wal_file")"
        else
            warn "WAL magic mismatch: $(basename "$wal_file") (got: ${magic})"
            wal_valid=false
        fi
    done

    # Restore DB (if present in backup)
    if [[ "$has_db" == "true" ]]; then
        info "Restoring DB..."
        mkdir -p "$TARGET_DB"
        rsync -a "${BACKUP_PATH}/db/" "${TARGET_DB}/"
        local db_files
        db_files=$(find "$TARGET_DB" -type f | wc -l)
        success "DB restored: ${db_files} files"
    else
        info "Backup is WAL-only, skipping DB restore"
    fi

    # Summary
    echo ""
    echo "=========================================="
    if [[ "$wal_valid" == "true" ]]; then
        echo -e "  ${GREEN}Restore complete${NC}"
    else
        echo -e "  ${YELLOW}Restore complete (with WAL warnings)${NC}"
    fi
    echo "=========================================="
    echo ""
    echo "  WAL: ${TARGET_WAL}"
    if [[ "$has_db" == "true" ]]; then
        echo "  DB:  ${TARGET_DB}"
    fi
    echo ""
    echo "Start StemeDB with:"
    echo "  STEMEDB_WAL_DIR=${TARGET_WAL} STEMEDB_DB_DIR=${TARGET_DB} cargo run --bin stemedb-api"
    echo ""
}

main "$@"