#!/usr/bin/env bash
# ============================================================================
# brain_restore_test.sh — Verify integrity of a Vincent AI brain backup
#
# Runs a battery of checks against a backup directory without modifying it.
#
# Usage:  brain_restore_test.sh <backup_directory>
#
# Checks:
#   1. Required files exist
#   2. events.db integrity via SQLite PRAGMA
#   3. Entity count in reasonable range
#   4. Python can import and query the restored DB
#   5. Manifest checksums are valid
#
# Exit code: 0 = all pass, 1 = one or more failures
# ============================================================================
set -euo pipefail

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

PASS_COUNT=0
FAIL_COUNT=0
SKIP_COUNT=0

pass() {
    PASS_COUNT=$((PASS_COUNT + 1))
    echo -e "  ${GREEN}✓ PASS${NC}  $1"
}

fail() {
    FAIL_COUNT=$((FAIL_COUNT + 1))
    echo -e "  ${RED}✗ FAIL${NC}  $1"
}

skip() {
    SKIP_COUNT=$((SKIP_COUNT + 1))
    echo -e "  ${YELLOW}⊘ SKIP${NC}  $1"
}

# ── Argument validation ──────────────────────────────────────────────────────
if [[ $# -lt 1 ]]; then
    echo "Usage: $0 <backup_directory>"
    echo "Example: $0 ~/backups/brain/2026-06-02"
    exit 1
fi

BACKUP_DIR="$1"

echo "═══════════════════════════════════════════════════════"
echo "  Brain Backup Restore Test"
echo "  Backup: $BACKUP_DIR"
echo "  Date:   $(date '+%Y-%m-%d %H:%M:%S')"
echo "═══════════════════════════════════════════════════════"
echo ""

# ── Pre-flight ───────────────────────────────────────────────────────────────
if [[ ! -d "$BACKUP_DIR" ]]; then
    fail "Backup directory does not exist: $BACKUP_DIR"
    echo ""
    echo "═══════════════════════════════════════════════════════"
    echo "  RESULT: $PASS_COUNT passed, $FAIL_COUNT failed, $SKIP_COUNT skipped"
    echo "═══════════════════════════════════════════════════════"
    exit 1
fi

# ── Test 1: Required files exist ─────────────────────────────────────────────
echo "--- Test 1: Required files present ---"

REQUIRED_FILES=("events.db" "config.py" "brain.py")
for f in "${REQUIRED_FILES[@]}"; do
    if [[ -f "$BACKUP_DIR/$f" ]]; then
        pass "$f exists ($(du -h "$BACKUP_DIR/$f" | cut -f1))"
    else
        fail "$f not found"
    fi
done

# Optional files (config extras)
OPTIONAL_FILES=("embeddings.py" "storage.py")
for f in "${OPTIONAL_FILES[@]}"; do
    if [[ -f "$BACKUP_DIR/$f" ]]; then
        pass "$f exists (optional)"
    else
        skip "$f not found (optional)"
    fi
done

# Directories
if [[ -d "$BACKUP_DIR/events" ]]; then
    event_files=$(find "$BACKUP_DIR/events" -type f -name '*.json' 2>/dev/null | wc -l)
    pass "events/ directory exists ($event_files JSON files)"
else
    skip "events/ directory not found"
fi

if [[ -d "$BACKUP_DIR/migrations" ]]; then
    mig_files=$(find "$BACKUP_DIR/migrations" -type f 2>/dev/null | wc -l)
    pass "migrations/ directory exists ($mig_files files)"
else
    skip "migrations/ directory not found"
fi

if [[ -f "$BACKUP_DIR/manifest.txt" ]]; then
    pass "manifest.txt exists"
else
    skip "manifest.txt not found"
fi

echo ""

# ── Test 2: SQLite integrity check ───────────────────────────────────────────
echo "--- Test 2: SQLite database integrity ---"

DB="$BACKUP_DIR/events.db"
if [[ -f "$DB" ]]; then
    integrity=$(sqlite3 "$DB" "PRAGMA integrity_check;" 2>&1)
    if [[ "$integrity" == "ok" ]]; then
        pass "PRAGMA integrity_check returned 'ok'"
    else
        fail "PRAGMA integrity_check returned: $integrity"
    fi

    # Check WAL mode status
    wal_mode=$(sqlite3 "$DB" "PRAGMA journal_mode;" 2>&1)
    if [[ $? -eq 0 ]]; then
        pass "Journal mode readable: $wal_mode"
    else
        fail "Cannot read journal mode"
    fi

    # Check free pages
    free_pages=$(sqlite3 "$DB" "PRAGMA freelist_count;" 2>&1)
    if [[ $? -eq 0 ]]; then
        pass "Free pages count: $free_pages"
    else
        fail "Cannot read freelist count"
    fi
else
    fail "events.db not found — cannot run integrity checks"
fi

echo ""

# ── Test 3: Entity count validation ──────────────────────────────────────────
echo "--- Test 3: Entity count in expected range ---"

if [[ -f "$DB" ]]; then
    # Count total rows across key tables
    total_entities=0
    tables_found=0

    # Check common brain schema tables
    for table in entities conversations contexts skills prompts; do
        exists=$(sqlite3 "$DB" "SELECT count(*) FROM sqlite_master WHERE type='table' AND name='$table';" 2>/dev/null || echo "0")
        if [[ "$exists" -gt 0 ]]; then
            count=$(sqlite3 "$DB" "SELECT count(*) FROM $table;" 2>/dev/null || echo "0")
            echo "   Table '$table': $count rows"
            total_entities=$((total_entities + count))
            tables_found=$((tables_found + 1))
        fi
    done

    if [[ $tables_found -gt 0 ]]; then
        pass "Found $tables_found tables, $total_entities total entities"

        # Sanity: should have at least some data in a real backup
        if [[ $total_entities -gt 0 ]]; then
            pass "Entity count ($total_entities) is above zero — data present"
        else
            fail "Entity count is 0 — backup may be empty"
        fi
    else
        # If no known tables, just check that the DB has any tables at all
        table_list=$(sqlite3 "$DB" ".tables" 2>/dev/null || echo "")
        if [[ -n "$table_list" ]]; then
            pass "Database has tables: $table_list"
        else
            skip "No recognizable tables in database"
        fi
    fi
else
    fail "events.db not found — cannot count entities"
fi

echo ""

# ── Test 4: Python import test ───────────────────────────────────────────────
echo "--- Test 4: Python can query restored database ---"

python_result=$(python3 -c "
import sqlite3, sys, os

db = sys.argv[1]
if not os.path.isfile(db):
    print('ERROR: database file not found')
    sys.exit(1)

try:
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    # Test basic read — list all tables
    cursor.execute(\"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;\")
    tables = [row[0] for row in cursor.fetchall()]

    # Test a SELECT on first table if one exists
    errors = []
    if tables:
        for tbl in tables[:3]:  # Check first 3 tables
            try:
                cursor.execute(f'SELECT count(*) FROM [{tbl}]')
                count = cursor.fetchone()[0]
                print(f'  OK: [{tbl}] has {count} rows')
            except Exception as e:
                errors.append(f'{tbl}: {e}')

    if errors:
        for e in errors:
            print(f'  WARN: {e}')

    conn.close()
    print('SUCCESS')
except Exception as e:
    print(f'ERROR: {e}')
    sys.exit(1)
" "$DB" 2>&1)

if echo "$python_result" | grep -q "SUCCESS"; then
    pass "Python successfully connected and queried the database"
    # Print table details indented
    echo "$python_result" | grep -E "^\s+(OK|WARN)" | sed 's/^/     /'
else
    fail "Python import/query test failed: $python_result"
fi

echo ""

# ── Test 5: Manifest checksum verification ───────────────────────────────────
echo "--- Test 5: Manifest checksum verification ---"

MANIFEST="$BACKUP_DIR/manifest.txt"
if [[ -f "$MANIFEST" ]]; then
    # Checksums in manifest are absolute paths; we need to verify from backup dir
    # Extract just the sha256sum lines and verify
    checksum_errors=0
    while IFS= read -r line; do
        # Skip comments and empty lines
        [[ "$line" =~ ^# ]] && continue
        [[ -z "$line" ]] && continue
        # Lines that don't look like sha256sum output: 64-char hex  filename
        if [[ "$line" =~ ^[a-f0-9]{64}\  ]]; then
            file=$(echo "$line" | awk '{print $2}')
            if [[ -f "$file" ]]; then
                expected=$(echo "$line" | awk '{print $1}')
                actual=$(sha256sum "$file" | awk '{print $1}')
                if [[ "$expected" == "$actual" ]]; then
                    : # checksum OK — no need to report each one
                else
                    fail "Checksum mismatch: $(basename "$file")"
                    checksum_errors=$((checksum_errors + 1))
                fi
            fi
        fi
    done < "$MANIFEST"

    if [[ $checksum_errors -eq 0 ]]; then
        # Count how many checksums we verified
        checksum_count=$(grep -cE '^[a-f0-9]{64}\ ' "$MANIFEST" 2>/dev/null || echo "0")
        pass "All $checksum_count checksums verified successfully"
    else
        fail "$checksum_errors checksum(s) failed verification"
    fi
else
    skip "manifest.txt not found — skipping checksum verification"
fi

echo ""

# ── Summary ──────────────────────────────────────────────────────────────────
echo "═══════════════════════════════════════════════════════"
if [[ $FAIL_COUNT -eq 0 ]]; then
    echo -e "  ${GREEN}ALL CHECKS PASSED${NC}"
else
    echo -e "  ${RED}$FAIL_COUNT CHECK(S) FAILED${NC}"
fi
echo "  Total: $PASS_COUNT passed, $FAIL_COUNT failed, $SKIP_COUNT skipped"
echo "═══════════════════════════════════════════════════════"

# Confirm backup was not modified
echo ""
echo "Backup directory was read-only during testing. No files modified."

exit $FAIL_COUNT