#!/bin/bash
#
# Start CyborgDB service with specified backing store for E2E tests.
#
# Usage:
#   ./start_service.sh <service_type> <backing_store> <image_tag>
#
# Arguments:
#   service_type: docker-x86, docker-arm, docker-gpu, standalone
#   backing_store: memory, redis, postgres, standalone
#   image_tag: Docker image tag (e.g., 0.16.0.dev20251211T142206)
#
# Environment variables:
#   CYBORGDB_API_KEY: API key for the service
#   PIP_EXTRA_INDEX_URL: CloudSmith PyPI URL (for standalone)
#
# Example:
#   ./start_service.sh docker-x86 redis 0.16.0.dev20251211T142206
#

set -e

SERVICE_TYPE="${1:-docker-x86}"
BACKING_STORE="${2:-memory}"
IMAGE_TAG="${3:-latest}"

# Detect OS
IS_MACOS=false
if [[ "$(uname)" == "Darwin" ]]; then
    IS_MACOS=true
fi

echo "Starting CyborgDB service..."
echo "  Service type: $SERVICE_TYPE"
echo "  Backing store: $BACKING_STORE"
echo "  Image tag: $IMAGE_TAG"
echo "  OS: $(uname) (macOS: $IS_MACOS)"

# Default environment variables
export CYBORGDB_HOST="${CYBORGDB_HOST:-0.0.0.0}"
export CYBORGDB_PORT="${CYBORGDB_PORT:-8000}"

# Configure backing store
configure_backing_store() {
    case "$BACKING_STORE" in
        memory)
            echo "Using memory backing store (valkeylite)"
            # Set INDEX_LOCATION, CONFIG_LOCATION, and ITEMS_LOCATION for the service settings validation
            export INDEX_LOCATION="memory"
            export CONFIG_LOCATION="memory"
            export ITEMS_LOCATION="memory"
            # Don't set CYBORGDB_DB_TYPE - service defaults to in-memory when not set
            # Only redis and postgres are supported values for CYBORGDB_DB_TYPE
            export CYBORGDB_DB_TYPE=""
            export CYBORGDB_CONNECTION_STRING=""
            ;;
        redis)
            if [[ "$IS_MACOS" == "true" ]]; then
                echo "Starting Redis via brew services (macOS)..."
                brew install redis 2>/dev/null || true
                brew services start redis
                # Wait for Redis to be ready
                for i in {1..30}; do
                    if redis-cli ping > /dev/null 2>&1; then
                        echo "Redis is ready"
                        break
                    fi
                    sleep 1
                done
            else
                echo "Starting Redis container..."
                docker run -d --name e2e-redis \
                    -p 6379:6379 \
                    redis:7-alpine
                sleep 2
            fi

            export INDEX_LOCATION="redis"
            export CONFIG_LOCATION="redis"
            export ITEMS_LOCATION="redis"
            export CYBORGDB_DB_TYPE="redis"
            export CYBORGDB_CONNECTION_STRING="host:localhost,port:6379,db:0"
            ;;
        postgres)
            if [[ "$IS_MACOS" == "true" ]]; then
                echo "Starting PostgreSQL via brew services (macOS)..."
                # GitHub's macos-15 runners have postgresql@14 pre-installed
                # Find and start the installed versioned postgresql formula
                PG_FORMULA=$(brew list 2>/dev/null | grep -E '^postgresql@[0-9]+$' | head -1)
                if [[ -z "$PG_FORMULA" ]]; then
                    echo "Installing postgresql@14..."
                    brew install postgresql@14 2>/dev/null || true
                    PG_FORMULA="postgresql@14"
                fi
                echo "Using $PG_FORMULA"

                # Get the Homebrew prefix for this formula
                PG_PREFIX=$(brew --prefix "$PG_FORMULA")
                echo "PostgreSQL prefix: $PG_PREFIX"

                # Ensure PostgreSQL binaries are in PATH
                export PATH="$PG_PREFIX/bin:$PATH"

                # Initialize database if needed (creates data directory)
                if [[ ! -d "$PG_PREFIX/var/$PG_FORMULA" ]]; then
                    echo "Initializing PostgreSQL database..."
                    initdb -D "$PG_PREFIX/var/$PG_FORMULA" 2>/dev/null || true
                fi

                # Start PostgreSQL
                brew services start "$PG_FORMULA"

                # Wait for PostgreSQL to be ready (try both TCP and socket)
                echo "Waiting for PostgreSQL to be ready..."
                for i in {1..30}; do
                    if pg_isready > /dev/null 2>&1; then
                        echo "PostgreSQL is ready"
                        break
                    fi
                    sleep 1
                done

                # On macOS Homebrew PostgreSQL, the default superuser is the current user (not 'postgres')
                # So we connect without -U flag to use the current user's credentials
                createdb cyborgdb 2>/dev/null || true
                psql -c "CREATE USER cyborgdb WITH PASSWORD 'testpassword';" 2>/dev/null || true
                psql -c "ALTER USER cyborgdb WITH SUPERUSER;" 2>/dev/null || true
                psql -d cyborgdb -c "GRANT ALL PRIVILEGES ON SCHEMA public TO cyborgdb;" 2>/dev/null || true

                # Use Unix socket connection (more reliable on macOS Homebrew)
                export INDEX_LOCATION="postgres"
                export CONFIG_LOCATION="postgres"
                export ITEMS_LOCATION="postgres"
                export CYBORGDB_DB_TYPE="postgres"
                export CYBORGDB_CONNECTION_STRING="host=/tmp dbname=cyborgdb user=cyborgdb password=testpassword"
            else
                echo "Starting PostgreSQL container..."
                docker run -d --name e2e-postgres \
                    -p 5432:5432 \
                    -e POSTGRES_DB=cyborgdb \
                    -e POSTGRES_USER=cyborgdb \
                    -e POSTGRES_PASSWORD=testpassword \
                    postgres:15-alpine

                echo "Waiting for PostgreSQL to be ready..."
                for i in {1..30}; do
                    if docker exec e2e-postgres pg_isready -U cyborgdb > /dev/null 2>&1; then
                        echo "PostgreSQL is ready"
                        break
                    fi
                    sleep 1
                done

                export INDEX_LOCATION="postgres"
                export CONFIG_LOCATION="postgres"
                export ITEMS_LOCATION="postgres"
                export CYBORGDB_DB_TYPE="postgres"
                export CYBORGDB_CONNECTION_STRING="host=localhost port=5432 dbname=cyborgdb user=cyborgdb password=testpassword"
            fi
            ;;
        standalone)
            echo "Using Standalone backing store"
            export INDEX_LOCATION="standalone"
            export CONFIG_LOCATION="standalone"
            export ITEMS_LOCATION="standalone"
            export CYBORGDB_DB_TYPE="standalone"
            # Standalone works without a connection string - uses default path
            export CYBORGDB_CONNECTION_STRING=""
            ;;
        minio)
            if [[ "$IS_MACOS" == "true" ]]; then
                echo "Installing MinIO via brew (macOS)..."
                brew install minio/stable/minio minio/stable/mc 2>/dev/null || true

                export MINIO_ROOT_USER=minioadmin
                export MINIO_ROOT_PASSWORD=minioadmin
                MINIO_DATA_DIR=$(mktemp -d)
                nohup minio server "$MINIO_DATA_DIR" --address ":9000" --console-address ":9001" > /tmp/minio.log 2>&1 &
                echo $! > /tmp/minio.pid
            else
                echo "Starting MinIO container..."
                docker run -d --name e2e-minio \
                    -p 9000:9000 \
                    -p 9001:9001 \
                    -e "MINIO_ROOT_USER=minioadmin" \
                    -e "MINIO_ROOT_PASSWORD=minioadmin" \
                    minio/minio server /data --console-address ":9001"
            fi

            echo "Waiting for MinIO to be ready..."
            for i in {1..30}; do
                if curl -sf http://localhost:9000/minio/health/ready > /dev/null 2>&1; then
                    echo "MinIO is ready"
                    break
                fi
                sleep 1
            done

            echo "Creating test bucket..."
            if [[ "$IS_MACOS" == "true" ]]; then
                mc alias set local http://localhost:9000 minioadmin minioadmin
                mc mb local/test-bucket
            else
                docker run --rm --network host \
                    --entrypoint sh minio/mc -c \
                    "mc alias set local http://localhost:9000 minioadmin minioadmin && mc mb local/test-bucket"
            fi

            export INDEX_LOCATION="s3"
            export CONFIG_LOCATION="s3"
            export ITEMS_LOCATION="s3"
            export CYBORGDB_DB_TYPE="s3"
            export CYBORGDB_CONNECTION_STRING="s3://test-bucket?region=us-east-1&endpoint=http://localhost:9000&access_key=minioadmin&secret_key=minioadmin"
            ;;
        *)
            echo "Error: Unknown backing store: $BACKING_STORE"
            exit 1
            ;;
    esac
}

# Start service based on type
start_service() {
    # Use DOCKER_REGISTRY env var if set, otherwise default to Cloudsmith
    local REGISTRY="${DOCKER_REGISTRY:-docker.cloudsmith.io/cyborg/cyborgdb-dev}"

    case "$SERVICE_TYPE" in
        docker-x86|docker-arm)
            IMAGE="${REGISTRY}/cyborgdb-service:${IMAGE_TAG}"

            echo "Starting Docker container: $IMAGE"
            echo "  DB_TYPE: ${CYBORGDB_DB_TYPE:-<not set>}"
            echo "  CONNECTION_STRING: ${CYBORGDB_CONNECTION_STRING:-<not set>}"

            # Build docker run command
            DOCKER_ARGS=(
                -d
                --name e2e-cyborgdb-service
                -p "${CYBORGDB_PORT}:8000"
                -e "CYBORGDB_API_KEY=${CYBORGDB_API_KEY}"
                -e "INDEX_LOCATION=${INDEX_LOCATION}"
                -e "CONFIG_LOCATION=${CONFIG_LOCATION}"
                -e "ITEMS_LOCATION=${ITEMS_LOCATION}"
            )

            # Only pass DB_TYPE if set (memory uses default in-memory store)
            if [ -n "$CYBORGDB_DB_TYPE" ]; then
                DOCKER_ARGS+=(-e "CYBORGDB_DB_TYPE=${CYBORGDB_DB_TYPE}")
            fi

            if [ -n "$CYBORGDB_CONNECTION_STRING" ]; then
                DOCKER_ARGS+=(-e "CYBORGDB_CONNECTION_STRING=${CYBORGDB_CONNECTION_STRING}")
            fi

            # Pass RETRAIN_THRESHOLD if set (disables auto-retraining when high)
            if [ -n "$RETRAIN_THRESHOLD" ]; then
                DOCKER_ARGS+=(-e "RETRAIN_THRESHOLD=${RETRAIN_THRESHOLD}")
            fi

            # Use host network for backing store containers (redis, postgres)
            if [[ "$BACKING_STORE" != "memory" ]]; then
                DOCKER_ARGS+=(--network host)
            fi

            docker run "${DOCKER_ARGS[@]}" "$IMAGE"
            ;;

        docker-gpu)
            IMAGE="${REGISTRY}/cyborgdb-service-cuda:${IMAGE_TAG}"
            echo "Starting GPU Docker container: $IMAGE"

            DOCKER_ARGS=(
                -d
                --name e2e-cyborgdb-service
                --gpus all
                -p "${CYBORGDB_PORT}:8000"
                -e "CYBORGDB_API_KEY=${CYBORGDB_API_KEY}"
                -e "INDEX_LOCATION=${INDEX_LOCATION}"
                -e "CONFIG_LOCATION=${CONFIG_LOCATION}"
                -e "ITEMS_LOCATION=${ITEMS_LOCATION}"
            )

            # Only pass DB_TYPE if set (memory uses default in-memory store)
            if [ -n "$CYBORGDB_DB_TYPE" ]; then
                DOCKER_ARGS+=(-e "CYBORGDB_DB_TYPE=${CYBORGDB_DB_TYPE}")
            fi

            if [ -n "$CYBORGDB_CONNECTION_STRING" ]; then
                DOCKER_ARGS+=(-e "CYBORGDB_CONNECTION_STRING=${CYBORGDB_CONNECTION_STRING}")
            fi

            # Pass RETRAIN_THRESHOLD if set (disables auto-retraining when high)
            if [ -n "$RETRAIN_THRESHOLD" ]; then
                DOCKER_ARGS+=(-e "RETRAIN_THRESHOLD=${RETRAIN_THRESHOLD}")
            fi

            if [[ "$BACKING_STORE" != "memory" ]]; then
                DOCKER_ARGS+=(--network host)
            fi

            docker run "${DOCKER_ARGS[@]}" "$IMAGE"
            ;;

        standalone)
            echo "Installing CyborgDB service standalone..."

            # Install the service package with embeddings support
            if [[ "${ENVIRONMENT}" == "prod" ]]; then
                pip install "cyborgdb-service[embeddings]"
            else
                # Use --pre to get latest dev version from Cloudsmith
                pip install "cyborgdb-service[embeddings]" \
                    --pre --extra-index-url "${PIP_EXTRA_INDEX_URL}"
            fi

            # Print installed version for debugging
            echo "Installed version: $(python -c 'import cyborgdb_service; print(cyborgdb_service.__version__)')"

            echo "Starting standalone service..."
            # Export all required environment variables for the service
            # Note: Export both with and without CYBORGDB_ prefix for compatibility
            export CYBORGDB_HOST="${CYBORGDB_HOST}"
            export CYBORGDB_PORT="${CYBORGDB_PORT}"
            export CYBORGDB_API_KEY="${CYBORGDB_API_KEY}"
            export INDEX_LOCATION="${INDEX_LOCATION}"
            export CONFIG_LOCATION="${CONFIG_LOCATION}"
            export ITEMS_LOCATION="${ITEMS_LOCATION}"
            # Also export with CYBORGDB_ prefix in case service expects that
            export CYBORGDB_INDEX_LOCATION="${INDEX_LOCATION}"
            export CYBORGDB_CONFIG_LOCATION="${CONFIG_LOCATION}"
            export CYBORGDB_ITEMS_LOCATION="${ITEMS_LOCATION}"
            # Only export DB_TYPE and CONNECTION_STRING if set
            if [ -n "$CYBORGDB_DB_TYPE" ]; then
                export CYBORGDB_DB_TYPE="${CYBORGDB_DB_TYPE}"
            fi
            if [ -n "$CYBORGDB_CONNECTION_STRING" ]; then
                export CYBORGDB_CONNECTION_STRING="${CYBORGDB_CONNECTION_STRING}"
            fi

            # Debug: print env vars
            echo "  INDEX_LOCATION=${INDEX_LOCATION}"
            echo "  CONFIG_LOCATION=${CONFIG_LOCATION}"
            echo "  ITEMS_LOCATION=${ITEMS_LOCATION}"
            echo "  CYBORGDB_DB_TYPE=${CYBORGDB_DB_TYPE}"
            echo "  CYBORGDB_CONNECTION_STRING=${CYBORGDB_CONNECTION_STRING}"

            # Start service in background with nohup to ensure it keeps running
            nohup cyborgdb-service > /tmp/cyborgdb-service.log 2>&1 &
            SERVICE_PID=$!
            echo "$SERVICE_PID" > /tmp/cyborgdb-service.pid
            echo "Service started with PID: $SERVICE_PID"
            sleep 5  # Give service time to start

            # Check if process is still running and show early logs
            if kill -0 "$SERVICE_PID" 2>/dev/null; then
                echo "Service process is running"
            else
                echo "WARNING: Service process died immediately!"
                echo "Service logs:"
                cat /tmp/cyborgdb-service.log 2>&1 || echo "No logs available"
            fi
            ;;

        *)
            echo "Error: Unknown service type: $SERVICE_TYPE"
            exit 1
            ;;
    esac
}

# Wait for service to be healthy
wait_for_health() {
    local max_attempts=60
    local attempt=1

    echo "Waiting for service to be healthy..."

    while [ $attempt -le $max_attempts ]; do
        if curl -sf "http://localhost:${CYBORGDB_PORT}/health" > /dev/null 2>&1; then
            echo "Service is healthy after $attempt attempts"
            return 0
        fi

        # Also try /v1/health endpoint
        if curl -sf "http://localhost:${CYBORGDB_PORT}/v1/health" > /dev/null 2>&1; then
            echo "Service is healthy after $attempt attempts"
            return 0
        fi

        echo "Attempt $attempt/$max_attempts - waiting..."
        sleep 2
        attempt=$((attempt + 1))
    done

    echo "Error: Service failed to become healthy after $max_attempts attempts"

    # Print logs for debugging
    if [[ "$SERVICE_TYPE" == docker* ]]; then
        echo "Container logs:"
        docker logs e2e-cyborgdb-service 2>&1 || true
    elif [[ "$SERVICE_TYPE" == "standalone" ]]; then
        echo "Standalone service logs:"
        cat /tmp/cyborgdb-service.log 2>&1 || echo "No service logs found"
    fi

    return 1
}

# Cleanup function
cleanup() {
    echo "Cleaning up E2E test environment..."

    # Stop service container
    docker rm -f e2e-cyborgdb-service 2>/dev/null || true

    # Stop backing store containers
    docker rm -f e2e-redis 2>/dev/null || true
    docker rm -f e2e-postgres 2>/dev/null || true
    docker rm -f e2e-minio 2>/dev/null || true

    # Stop MinIO process (macOS)
    if [ -f /tmp/minio.pid ]; then
        kill "$(cat /tmp/minio.pid)" 2>/dev/null || true
        rm /tmp/minio.pid
    fi

    # Stop standalone service
    if [ -f /tmp/cyborgdb-service.pid ]; then
        kill "$(cat /tmp/cyborgdb-service.pid)" 2>/dev/null || true
        rm /tmp/cyborgdb-service.pid
    fi

    echo "Cleanup complete"
}

# Handle cleanup on exit if --cleanup flag is passed
if [[ "$4" == "--cleanup" ]]; then
    cleanup
    exit 0
fi

# Main execution
configure_backing_store
start_service
wait_for_health

echo ""
echo "CyborgDB service is ready!"
echo "  URL: http://localhost:${CYBORGDB_PORT}"
echo "  Backing store: $BACKING_STORE"
echo ""
