#!/bin/bash
#==============================================================================
# Momentry Core — Upgrade Script
# Usage: bash upgrade_momentry.sh <delivery_dir>
#   <delivery_dir> : path to delivery package (e.g. release/delivery/v1.0.0_xxx/)
#
# Upgrades an existing momentry_core installation from a delivery package.
# 1. Pre-flight: version check, backup, health
# 2. Migration: apply DB schema changes
# 3. Binary: replace momentry binary
# 4. Scripts: replace Python processor scripts
# 5. Verification: health check + API smoke test
#==============================================================================

set -euo pipefail

PROJECT_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
COLOR=true

# ─── Color helpers ───
if [ "$COLOR" = true ]; then
    R='\033[0;31m'; G='\033[0;32m'; Y='\033[1;33m'; B='\033[0;34m'; C='\033[0;36m'; N='\033[0m'
else
    R=''; G=''; Y=''; B=''; C=''; N=''
fi
ok()  { echo -e "  ${G}✓${N} $1"; }
fail() { echo -e "  ${R}✗${N} $1"; FAILURES+=("$1"); }
info() { echo -e "  ${B}→${N} $1"; }
warn() { echo -e "  ${Y}⚠${N} $1"; }
header() { echo -e "\n${C}─── $1 ───${N}"; }

FAILURES=()
DELIVERY_DIR=""

# ─── Parse args ───
while [ $# -gt 0 ]; do
    case "$1" in
        --project-dir) PROJECT_DIR="$2"; shift 2 ;;
        --no-color) COLOR=false; shift ;;
        --help) head -20 "$0"; exit 0 ;;
        *)
            [ -z "$DELIVERY_DIR" ] && DELIVERY_DIR="$1" || {
                echo "Unknown: $1"; exit 1; }
            shift ;;
    esac
done

if [ -z "$DELIVERY_DIR" ]; then
    echo "Usage: bash upgrade_momentry.sh <delivery_dir>"
    echo ""
    echo "  <delivery_dir>  Path to delivery package directory"
    echo "                  e.g. $PROJECT_DIR/release/delivery/v1.0.0_0e73d2a_20260515_084750"
    echo ""
    echo "  Delivery package must contain:"
    echo "    - momentry_v*           (production binary)"
    echo "    - scripts/              (processor scripts)"
    echo "    - migrate_*.sql         (DB migrations)"
    echo "    - INSTALL.md or similar (package info)"
    exit 1
fi

DELIVERY_DIR="$(cd "$DELIVERY_DIR" 2>/dev/null && pwd)" || {
    echo "ERROR: Delivery directory not found: $DELIVERY_DIR"
    exit 1
}

PG_BIN="${PG_BIN:-$HOME/pgsql/18.3/bin}"
DB_NAME="${DB_NAME:-momentry}"
BACKUP_DIR="$PROJECT_DIR/backup/upgrade_$(date +%Y%m%d_%H%M%S)"

cd "$PROJECT_DIR"
echo -e "${C}========================================${N}"
echo -e "${C}  Momentry Core — Upgrade${N}"
echo -e "${C}  Delivery: $DELIVERY_DIR${N}"
echo -e "${C}  Project:  $PROJECT_DIR${N}"
echo -e "${C}  Date:     $(date '+%Y-%m-%d %H:%M:%S')${N}"
echo -e "${C}========================================${N}"

# ═══════════════════════════════════════════════════════════
# Phase 0: Pre-flight
# ═══════════════════════════════════════════════════════════
header "Phase 0/5 — Pre-flight"

# 0a. Check delivery package integrity
info "Checking delivery package..."
REQUIRED=()
BINARY=$(ls "$DELIVERY_DIR"/momentry_v* 2>/dev/null | head -1 || true)
if [ -n "$BINARY" ]; then
    ok "Binary: $(basename "$BINARY") ($(ls -lh "$BINARY" | awk '{print $5}'))"
else
    fail "No momentry_v* binary in delivery package"
fi

if [ -d "$DELIVERY_DIR/scripts" ]; then
    SCRIPT_COUNT=$(find "$DELIVERY_DIR/scripts" -name '*.py' -type f | wc -l | tr -d ' ')
    ok "Scripts dir: $SCRIPT_COUNT .py files"
else
    fail "No scripts/ directory in delivery package"
fi

MIGRATIONS=()
for f in "$DELIVERY_DIR"/migrate_*.sql; do
    [ -f "$f" ] && MIGRATIONS+=("$f")
done
if [ ${#MIGRATIONS[@]} -gt 0 ]; then
    for f in "${MIGRATIONS[@]}"; do ok "Migration: $(basename "$f")"; done
else
    warn "No migration SQL files in delivery package"
fi

# 0b. Check current server
HEALTH_ENDPOINT="http://127.0.0.1:3003"
CURL_OPTS="-sf --connect-timeout 5"

if curl $CURL_OPTS "$HEALTH_ENDPOINT/health" &>/dev/null; then
    CURRENT_VER=$(curl -sf "$HEALTH_ENDPOINT/health" | python3 -c "import json,sys;print(json.load(sys.stdin).get('version','?'))" 2>/dev/null)
    CURRENT_HASH=$(curl -sf "$HEALTH_ENDPOINT/health" | python3 -c "import json,sys;print(json.load(sys.stdin).get('build_git_hash','?'))" 2>/dev/null)
    ok "Server running: v$CURRENT_VER ($CURRENT_HASH)"
else
    CURRENT_VER="down"
    warn "Server not reachable at $HEALTH_ENDPOINT"
    HEALTH_ENDPOINT="http://127.0.0.1:3002"
    if curl $CURL_OPTS "$HEALTH_ENDPOINT/health" &>/dev/null; then
        CURRENT_VER=$(curl -sf "$HEALTH_ENDPOINT/health" | python3 -c "import json,sys;print(json.load(sys.stdin).get('version','?'))" 2>/dev/null)
        ok "Production server running: v$CURRENT_VER"
    fi
fi

# Extract package version from INSTALL.md or binary name
PKG_VER=$(echo "$(basename "$DELIVERY_DIR")" | grep -oE 'v[0-9]+\.[0-9]+\.[0-9]+' || echo "unknown")
PKG_BUILD=$(echo "$(basename "$DELIVERY_DIR")" | grep -oE '[a-f0-9]{7}' || echo "unknown")
info "Package: $PKG_VER (build $PKG_BUILD)"

# 0c. Check PostgreSQL
if "$PG_BIN/pg_isready" -q 2>/dev/null; then
    ok "PostgreSQL"
else
    fail "PostgreSQL not running"
fi

# 0d. Prompt to continue
if [ ${#FAILURES[@]} -gt 0 ]; then
    echo ""
    echo -e "${R}Pre-flight failures detected. Continue anyway? [y/N]${N} "
    read -r CONFIRM
    if [ "$CONFIRM" != "y" ] && [ "$CONFIRM" != "Y" ]; then
        echo "Aborted."
        exit 1
    fi
fi

# ═══════════════════════════════════════════════════════════
# Phase 1: Backup
# ═══════════════════════════════════════════════════════════
header "Phase 1/5 — Backup"

mkdir -p "$BACKUP_DIR"

# 1a. Backup current binary (if exists)
for bin_path in target/release/momentry target/debug/momentry_playground; do
    if [ -f "$PROJECT_DIR/$bin_path" ]; then
        mkdir -p "$BACKUP_DIR/$(dirname "$bin_path")"
        cp "$PROJECT_DIR/$bin_path" "$BACKUP_DIR/$bin_path"
        ok "Backed up: $bin_path"
    fi
done

# 1b. Backup scripts
SCRIPTS_SNAPSHOT="$BACKUP_DIR/scripts"
mkdir -p "$SCRIPTS_SNAPSHOT"
rsync -a "$PROJECT_DIR/scripts/" "$SCRIPTS_SNAPSHOT/" 2>/dev/null
ok "Backed up: scripts/ (to $SCRIPTS_SNAPSHOT)"

# 1c. Backup current .env
if [ -f "$PROJECT_DIR/.env" ]; then
    cp "$PROJECT_DIR/.env" "$BACKUP_DIR/.env"
    ok "Backed up: .env"
fi

# 1d. Schema backup
"$PG_BIN/pg_dump" -U accusys -d "$DB_NAME" --schema-only > "$BACKUP_DIR/schema_pre.sql" 2>/dev/null
ok "Backed up: schema (pre-upgrade)"

echo ""
info "Backup saved to: $BACKUP_DIR"

# ═══════════════════════════════════════════════════════════
# Phase 2: Database Migration
# ═══════════════════════════════════════════════════════════
header "Phase 2/6 — Database Migration"

# Ensure schema_migrations table exists first
"$PG_BIN/psql" -U accusys -d "$DB_NAME" -c "
    CREATE TABLE IF NOT EXISTS schema_migrations (
        id SERIAL PRIMARY KEY,
        filename TEXT NOT NULL UNIQUE,
        checksum TEXT NOT NULL,
        applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
        duration_ms INTEGER DEFAULT 0
    )" &>/dev/null
ok "schema_migrations table ensured"

for mig in "${MIGRATIONS[@]}"; do
    MIG_NAME=$(basename "$mig")
    MIG_HASH=$(shasum -a 256 "$mig" | awk '{print $1}')

    # Check if already applied with matching checksum
    ALREADY=$("$PG_BIN/psql" -U accusys -d "$DB_NAME" -t -A -c \
        "SELECT COUNT(*) FROM schema_migrations WHERE filename='$MIG_NAME' AND checksum='$MIG_HASH'" 2>/dev/null || echo "0")
    if [ "$ALREADY" -gt 0 ]; then
        ok "Migration $MIG_NAME (already applied, checksum match)"
        continue
    fi

    info "Applying $MIG_NAME..."
    T0=$(date +%s%N)
    if "$PG_BIN/psql" -U accusys -d "$DB_NAME" -f "$mig" 2>&1; then
        T1=$(date +%s%N)
        DURATION_MS=$(( (T1 - T0) / 1000000 ))
        "$PG_BIN/psql" -U accusys -d "$DB_NAME" -c \
            "INSERT INTO schema_migrations (filename, checksum, duration_ms) VALUES ('$MIG_NAME', '$MIG_HASH', $DURATION_MS) ON CONFLICT (filename) DO UPDATE SET checksum=EXCLUDED.checksum, applied_at=NOW()" &>/dev/null || true
        ok "Applied: $MIG_NAME (${DURATION_MS}ms)"
    else
        fail "Migration failed: $MIG_NAME"
    fi
done

# Verify schema completeness via binary startup check
info "Verifying schema... (will be checked at server restart)"

# ═══════════════════════════════════════════════════════════
# Phase 3: Binary Replacement
# ═══════════════════════════════════════════════════════════
header "Phase 3/5 — Binary Replacement"

# Determine which binary to install
if echo "$HEALTH_ENDPOINT" | grep -q "3002"; then
    TARGET_BIN="momentry"
    TARGET_PATH="$PROJECT_DIR/target/release/$TARGET_BIN"
    TARGET_PORT=3002
else
    TARGET_BIN="momentry_playground"
    TARGET_PATH="$PROJECT_DIR/target/debug/$TARGET_BIN"
    TARGET_PORT=3003
fi

# Copy binary from delivery
BINARY_SRC=$(ls "$DELIVERY_DIR"/momentry_v* 2>/dev/null | head -1 || true)
if [ -n "$BINARY_SRC" ]; then
    # If the delivery has a pre-built binary (for production)
    mkdir -p "$(dirname "$TARGET_PATH")"
    cp "$BINARY_SRC" "$TARGET_PATH"
    chmod +x "$TARGET_PATH"

    # Remove code signature (required for Apple Silicon when binary is from another Mac)
    if command -v codesign &>/dev/null; then
        codesign --remove-signature "$TARGET_PATH" 2>/dev/null || true
    fi

    ok "Binary replaced: $TARGET_PATH ($(ls -lh "$TARGET_PATH" | awk '{print $5}'))"
else
    # No pre-built binary — need to build from source
    info "No pre-built binary in delivery — building from source..."
    cargo build --bin "$TARGET_BIN" 2>&1 | tail -3
    if [ -f "$PROJECT_DIR/target/debug/momentry_playground" ]; then
        ok "Built: $TARGET_BIN"
    else
        fail "Build failed: $TARGET_BIN"
    fi
fi

# ═══════════════════════════════════════════════════════════
# Phase 4: Scripts Update
# ═══════════════════════════════════════════════════════════
header "Phase 4/5 — Scripts Update"

if [ -d "$DELIVERY_DIR/scripts" ]; then
    # Copy each processor script, preserving existing utility scripts
    # that may not be in the delivery package
    for f in "$DELIVERY_DIR/scripts"/*.py; do
        if [ -f "$f" ]; then
            cp "$f" "$PROJECT_DIR/scripts/"
        fi
    done
    # Copy subdirectories (swift_processors, utils, etc.)
    for d in "$DELIVERY_DIR/scripts"/*/; do
        if [ -d "$d" ]; then
            rsync -a "$d" "$PROJECT_DIR/scripts/$(basename "$d")/"
        fi
    done
    chmod +x "$PROJECT_DIR/scripts"/*.py 2>/dev/null || true
    ok "Processor scripts updated ($(find "$PROJECT_DIR/scripts" -name '*.py' -type f | wc -l | tr -d ' ') .py files)"

    # Verify SHA256 checksums after update
    CHECKSUMS_FILE="$PROJECT_DIR/scripts/checksums.sha256"
    CS_TOTAL=0; CS_PASS=0; CS_FAIL=0
    if [ -f "$CHECKSUMS_FILE" ]; then
        while IFS= read -r line; do
            [ -z "$line" ] && continue
            EXPECTED_HASH=$(echo "$line" | awk '{print $1}')
            FILE_PATH=$(echo "$line" | awk '{print $2}')
            FULL_PATH="$PROJECT_DIR/scripts/$FILE_PATH"
            CS_TOTAL=$((CS_TOTAL + 1))
            if [ -f "$FULL_PATH" ]; then
                ACTUAL_HASH=$(shasum -a 256 "$FULL_PATH" 2>/dev/null | awk '{print $1}')
                if [ "$ACTUAL_HASH" = "$EXPECTED_HASH" ]; then
                    CS_PASS=$((CS_PASS + 1))
                else
                    CS_FAIL=$((CS_FAIL + 1))
                    [ $CS_FAIL -le 5 ] && warn "$FILE_PATH — hash mismatch after update"
                fi
            else
                CS_FAIL=$((CS_FAIL + 1))
                [ $CS_FAIL -le 5 ] && warn "$FILE_PATH — not found after update"
            fi
        done < "$CHECKSUMS_FILE"
        ok "Script integrity: $CS_PASS/$CS_TOTAL checksums match"
        [ $CS_FAIL -gt 0 ] && warn "Script integrity: $CS_FAIL mismatches" || true
    else
        warn "No checksums.sha256 found — generating from updated scripts..."
        cd "$PROJECT_DIR/scripts"
        find . -name '*.py' -type f | sort | while IFS= read -r f; do shasum -a 256 "$f"; done > checksums.sha256
        ok "checksums.sha256 generated ($(wc -l < checksums.sha256) entries)"
    fi
else
    warn "No scripts/ in delivery — skipping"
fi

# Python dependencies check
header "  Python dependencies"
PYTHON_BIN="${MOMENTRY_PYTHON_PATH:-/opt/homebrew/bin/python3.11}"
for pkg in PyPDF2 python-docx openpyxl python-pptx; do
    if "$PYTHON_BIN" -c "import ${pkg%%=*}" &>/dev/null 2>&1; then
        ok "$pkg"
    else
        "$PYTHON_BIN" -m pip install "$pkg" --quiet && ok "$pkg" || warn "$pkg"
    fi
done

# Ensure watch directory exists
WATCH_DIR="${MOMENTRY_SFTP_ROOT:-$PROJECT_DIR/storage/watch}"
mkdir -p "$WATCH_DIR"
ok "Watcher directory: $WATCH_DIR"

# ═══════════════════════════════════════════════════════════
# Phase 6: Restart & Verify
# ═══════════════════════════════════════════════════════════
header "Phase 6/6 — Restart & Verify"

# 5a. Stop old server
if [ -n "$TARGET_PORT" ] && [ "$TARGET_PORT" = "3002" ]; then
    pkill -f "momentry server" 2>/dev/null || true
    pkill -f "momentry_playground server" 2>/dev/null || true
else
    pkill -f "momentry_playground server" 2>/dev/null || true
fi
sleep 2
info "Stopped old server"

# 5b. Start new server
cd "$PROJECT_DIR"
if [ "$TARGET_PORT" = "3002" ]; then
    DATABASE_SCHEMA=public nohup "$TARGET_PATH" server --port 3002 > "$PROJECT_DIR/production_boot.log" 2>&1 &
else
    DATABASE_SCHEMA=dev nohup "$TARGET_PATH" server --port 3003 > "$PROJECT_DIR/playground_boot.log" 2>&1 &
fi

for i in $(seq 1 10); do
    sleep 2
    if curl $CURL_OPTS "$HEALTH_ENDPOINT/health" &>/dev/null; then
        ok "Server restarted on port $TARGET_PORT"
        break
    fi
done
if ! curl $CURL_OPTS "$HEALTH_ENDPOINT/health" &>/dev/null; then
    fail "Server failed to start on port $TARGET_PORT"
fi

# 5c. Health check
HEALTH=$(curl -sf "$HEALTH_ENDPOINT/health" 2>/dev/null || echo '{"status":"error"}')
echo "$HEALTH" | python3 -c "
import json,sys
d=json.load(sys.stdin)
print(f'    Version:   {d.get(\"version\",\"?\")}')
print(f'    Build:     {d.get(\"build_git_hash\",\"?\")}')
print(f'    Timestamp: {d.get(\"build_timestamp\",\"?\")}')
print(f'    Status:    {d.get(\"status\",\"?\")}')
" 2>/dev/null

# 5d. Processor inventory
DETAILED=$(curl -sf "$HEALTH_ENDPOINT/health/detailed" 2>/dev/null || echo '{}')
echo "$DETAILED" | python3 -c "
import json,sys
d=json.load(sys.stdin)
p=d.get('pipeline',{})
proc=p.get('processors',{})
missing=[k for k in ['asr','yolo','face','pose','ocr','cut','caption','scene','story','asrx','probe','visual_chunk'] if not proc.get(k)]
if missing:
    print(f'    Missing processors: {missing}')
else:
    print('    All 12 processors: ✓')
print(f'    Scripts: {p.get(\"scripts_count\",\"?\")} .py files')
" 2>/dev/null

# 5e. API smoke test
API_KEY="${MOMENTRY_API_KEY:-muser_68600856036340bcafc01930eb4bd839_1774418104_97221b69}"
STATUS=$(curl -sf -o /dev/null -w "%{http_code}" \
    -H "X-API-Key: $API_KEY" \
    "$HEALTH_ENDPOINT/api/v1/videos?page=1&page_size=1" 2>/dev/null || echo "000")
ok "API /v1/videos → HTTP $STATUS"

# ═══════════════════════════════════════════════════════════
# Summary
# ═══════════════════════════════════════════════════════════
echo ""
echo -e "${C}========================================${N}"
if [ ${#FAILURES[@]} -eq 0 ]; then
    echo -e "${G}  Upgrade Complete — All Checks Passed${N}"
else
    echo -e "${Y}  Upgrade Complete — ${#FAILURES[@]} Warnings${N}"
    for f in "${FAILURES[@]}"; do echo -e "  ${R}✗${N} $f"; done
fi
echo -e "${C}========================================${N}"
echo ""
echo "  Backup:     $BACKUP_DIR"
echo "  Server:     $HEALTH_ENDPOINT"
echo "  Binary:     $TARGET_PATH"
echo "  Watch dir:  $WATCH_DIR"
echo "  Version:    $(echo "$HEALTH" | python3 -c "import json,sys;print(json.load(sys.stdin).get('version','?'))" 2>/dev/null)"
echo ""
echo "  To rollback:"
echo "    1. cp $BACKUP_DIR/target/release/momentry $PROJECT_DIR/target/release/"
echo "    2. rsync -a $BACKUP_DIR/scripts/ $PROJECT_DIR/scripts/"
echo "    3. $PG_BIN/psql -U accusys -d $DB_NAME < $BACKUP_DIR/schema_pre.sql"
echo "    4. Restart server"
echo ""

exit $([ ${#FAILURES[@]} -eq 0 ] && echo 0 || echo 1)