#!/bin/bash
set -e

SNAPSHOT_DIR="/data/qdrant_snapshots"
LOCAL_SNAPSHOT_DIR="/tmp/qdrant_snapshots"
QDRANT_URL="http://localhost:6333"
API_KEY="${QDRANT__SERVICE__API_KEY}"

# Copy snapshots from FUSE mount to local disk first
echo "Copying snapshots from FUSE mount to local /tmp..."
mkdir -p "$LOCAL_SNAPSHOT_DIR"
cp -r "$SNAPSHOT_DIR"/. "$LOCAL_SNAPSHOT_DIR/"
echo "Copy complete."

# Start Qdrant in background
/qdrant/qdrant &
QDRANT_PID=$!

# Wait for Qdrant to be ready
echo "Waiting for Qdrant to start..."
until wget -qO- --header="api-key: $API_KEY" "$QDRANT_URL/readyz" > /dev/null 2>&1; do
    sleep 1
done
echo "Qdrant ready."

# Recover snapshots from LOCAL copy (not FUSE mount)
echo "Scanning $LOCAL_SNAPSHOT_DIR for snapshots..."

for COLLECTION_DIR in "$LOCAL_SNAPSHOT_DIR"/*/; do
    [ -d "$COLLECTION_DIR" ] || continue
    COLLECTION=$(basename "$COLLECTION_DIR")
    echo "Found collection folder: $COLLECTION"

    SNAPSHOT=""
    for f in "$COLLECTION_DIR"*.snapshot; do
        case "$f" in *.checksum) continue ;; esac
        [ -f "$f" ] && SNAPSHOT="$f" && break
    done

    if [ -z "$SNAPSHOT" ]; then
        echo "No snapshot file found in $COLLECTION_DIR, skipping."
        continue
    fi
    echo "Found snapshot: $SNAPSHOT"

    HTTP_STATUS=$(wget -qO- \
        --header="api-key: $API_KEY" \
        "$QDRANT_URL/collections/$COLLECTION" 2>/dev/null | grep -o '"status":"[^"]*"' | head -1)
    echo "Collection $COLLECTION status: $HTTP_STATUS"

    if echo "$HTTP_STATUS" | grep -q '"status":"green"'; then
        echo "Collection '$COLLECTION' already green, skipping."
        continue
    fi

    echo "Recovering '$COLLECTION' from $SNAPSHOT ..."
    RESPONSE=$(wget -qO- \
        --method=PUT \
        --header="Content-Type: application/json" \
        --header="api-key: $API_KEY" \
        --body-data="{\"location\":\"file://$SNAPSHOT\"}" \
        "$QDRANT_URL/collections/$COLLECTION/snapshots/recover" 2>&1)
    echo "Recovery response: $RESPONSE"

    echo "Waiting for '$COLLECTION' to become green..."
    for i in $(seq 1 60); do
        sleep 5
        STATUS=$(wget -qO- \
            --header="api-key: $API_KEY" \
            "$QDRANT_URL/collections/$COLLECTION" 2>/dev/null | grep -o '"status":"[^"]*"' | head -1)
        echo "  Attempt $i: $STATUS"
        if echo "$STATUS" | grep -q '"status":"green"'; then
            echo "  '$COLLECTION' is green!"
            break
        fi
    done
done

echo "Snapshot recovery complete."
wait $QDRANT_PID