#!/bin/bash
# Auto-restart wrapper for dataset generation
# Monitors and restarts until target is reached

OUTPUT="${1:-./data/dataset_60k.pt}"
TARGET="${2:-60000}"
LOG_DIR="$(dirname $OUTPUT)"

echo "=============================================="
echo "Auto-restart Dataset Generator"
echo "Output: $OUTPUT"
echo "Target: $TARGET items"
echo "=============================================="

# Load environment
if [ -f .env ]; then
    export $(grep -v '^#' .env | xargs)
fi

get_count() {
    # Count = base (output file) + new (checkpoint.new)
    local new_checkpoint="${OUTPUT}.checkpoint.new"
    local base_count=0
    local new_count=0

    if [ -f "$OUTPUT" ]; then
        base_count=$(python3 -c "import torch; d=torch.load('$OUTPUT', map_location='cpu', weights_only=False, mmap=True); print(len(d))" 2>/dev/null || echo "0")
    fi

    if [ -f "$new_checkpoint" ]; then
        new_count=$(python3 -c "import torch; d=torch.load('$new_checkpoint', map_location='cpu', weights_only=False, mmap=True); print(len(d))" 2>/dev/null || echo "0")
    fi

    echo $((base_count + new_count))
}

attempt=1
while true; do
    current=$(get_count)
    echo ""
    echo "[Attempt $attempt] Current: $current / $TARGET items"

    if [ "$current" -ge "$TARGET" ]; then
        echo "=============================================="
        echo "COMPLETE! $current items generated"
        echo "=============================================="
        exit 0
    fi

    remaining=$((TARGET - current))
    echo "[Attempt $attempt] Starting generation (need $remaining more)..."
    echo ""

    # Run the generator with resume flag
    python3 datasets/create_dataset.py \
        --output "$OUTPUT" \
        --count "$TARGET" \
        --resume \
        2>&1 | tee -a "${LOG_DIR}/auto_restart.log"

    exit_code=$?

    # Check result
    new_count=$(get_count)
    echo ""
    echo "[Attempt $attempt] Finished with exit code $exit_code"
    echo "[Attempt $attempt] Items after run: $new_count"

    if [ "$new_count" -ge "$TARGET" ]; then
        echo "=============================================="
        echo "COMPLETE! $new_count items generated"
        echo "=============================================="
        exit 0
    fi

    # If no progress was made, wait longer before retry
    if [ "$new_count" -eq "$current" ]; then
        echo "[WARN] No progress made, waiting 30s before retry..."
        sleep 30
    else
        echo "[INFO] Progress: $current -> $new_count (+$((new_count - current)))"
        echo "[INFO] Restarting in 5s..."
        sleep 5
    fi

    attempt=$((attempt + 1))

    # Safety limit
    if [ "$attempt" -gt 50 ]; then
        echo "[ERROR] Too many attempts (50), giving up"
        exit 1
    fi
done