#!/bin/bash

set -e

COMPOSE_PROJECT_NAME="${COMPOSE_PROJECT_NAME:-mem0}"
VOLUMES=("qdrant_data" "neo4j_data" "neo4j_logs" "neo4j_import" "neo4j_plugins")

echo "=========================================="
echo "   Mem0 Setup Script"
echo "=========================================="
echo ""

check_volumes_exist() {
    local exists=false
    for vol in "${VOLUMES[@]}"; do
        full_name="${COMPOSE_PROJECT_NAME}_${vol}"
        if docker volume ls -q | grep -q "^${full_name}$"; then
            exists=true
            break
        fi
    done
    echo "$exists"
}

list_existing_volumes() {
    echo "Existing volumes:"
    for vol in "${VOLUMES[@]}"; do
        full_name="${COMPOSE_PROJECT_NAME}_${vol}"
        if docker volume ls -q | grep -q "^${full_name}$"; then
            size=$(docker system df -v 2>/dev/null | grep "$full_name" | awk '{print $4}' || echo "unknown")
            echo "  - $full_name ($size)"
        fi
    done
}

remove_volumes() {
    echo "Stopping containers..."
    docker compose down 2>/dev/null || true
    
    echo "Removing volumes..."
    for vol in "${VOLUMES[@]}"; do
        full_name="${COMPOSE_PROJECT_NAME}_${vol}"
        if docker volume ls -q | grep -q "^${full_name}$"; then
            echo "  Removing $full_name..."
            docker volume rm "$full_name" 2>/dev/null || true
        fi
    done
    echo "Volumes removed."
}

build_and_start() {
    echo ""
    echo "Building containers..."
    docker compose build
    
    echo ""
    echo "Starting services..."
    docker compose up -d
    
    echo ""
    echo "Waiting for services to be healthy..."
    sleep 5
    
    echo ""
    echo "Checking health..."
    curl -s http://localhost:8000/health 2>/dev/null | jq . || echo "Health check not available yet. Services may still be starting."
    
    echo ""
    echo "=========================================="
    echo "   Setup Complete!"
    echo "=========================================="
    echo ""
    echo "Services:"
    echo "  - Backend API: http://localhost:8000"
    echo "  - API Docs: http://localhost:8000/docs"
    echo "  - Health Check: http://localhost:8000/health"
    echo ""
    echo "Logs: docker compose logs -f backend"
}

if [ "$(check_volumes_exist)" = "true" ]; then
    echo "Existing data volumes detected!"
    echo ""
    list_existing_volumes
    echo ""
    echo "Options:"
    echo "  1) Keep existing data and start services"
    echo "  2) Reset everything (DELETE ALL DATA) and start fresh"
    echo "  3) Exit"
    echo ""
    read -p "Choose an option [1/2/3]: " choice
    
    case $choice in
        1)
            echo ""
            echo "Keeping existing data..."
            build_and_start
            ;;
        2)
            echo ""
            read -p "Are you sure you want to DELETE ALL DATA? Type 'yes' to confirm: " confirm
            if [ "$confirm" = "yes" ]; then
                remove_volumes
                build_and_start
            else
                echo "Aborted."
                exit 1
            fi
            ;;
        3)
            echo "Exiting."
            exit 0
            ;;
        *)
            echo "Invalid option. Exiting."
            exit 1
            ;;
    esac
else
    echo "No existing volumes found. Starting fresh setup..."
    build_and_start
fi