#!/bin/bash
# prepare_environment.sh
#
# This script prepares the environment for running the Synthea pipeline.
# It rebuilds the Docker containers, ensures the directory structure is correct,
# and starts the necessary services.

set -e  # Exit on error

# Display a header
echo "=================================================="
echo "Synthea All Diseases Pipeline - Environment Setup"
echo "=================================================="

# Check for Docker
if ! command -v docker &> /dev/null; then
    echo "ERROR: Docker is not installed or not in the PATH."
    echo "Please install Docker and try again."
    exit 1
fi

# Check for docker-compose
if ! command -v docker-compose &> /dev/null; then
    echo "ERROR: docker-compose is not installed or not in the PATH."
    echo "Please install docker-compose and try again."
    exit 1
fi

# Create necessary directories if they don't exist
echo "Creating required directories..."
mkdir -p module_generator
mkdir -p src/main/resources/modules
mkdir -p src/main/resources/disease
mkdir -p modules
mkdir -p output

# Check if the module_generator directory has the required scripts
if [ ! -f "module_generator/module_generator.py" ] || [ ! -f "module_generator/run_module_generator.py" ]; then
    echo "ERROR: Required Python scripts not found in module_generator directory!"
    # Check if they're in the python directory and we need to move them
    if [ -d "python" ] && [ -f "python/module_generator.py" ] && [ -f "python/run_module_generator.py" ]; then
        echo "Found scripts in python directory, moving them to module_generator directory..."
        mkdir -p module_generator
        mv python/module_generator.py module_generator/
        mv python/run_module_generator.py module_generator/
        # Copy README if exists
        if [ -f "python/README_module_generator.md" ]; then
            mv python/README_module_generator.md module_generator/
        fi
    else
        echo "Scripts not found in python directory either. Please ensure the scripts are in the module_generator directory."
        exit 1
    fi
fi

# Make the Python scripts executable
chmod +x module_generator/module_generator.py
chmod +x module_generator/run_module_generator.py

# Create symlinks for convenience
ln -sf "$(pwd)/module_generator/run_module_generator.py" "$(pwd)/run_module_generator.py"
ln -sf "$(pwd)/module_generator/module_generator.py" "$(pwd)/module_generator.py"

# Create an empty disease_list.json if it doesn't exist
if [ ! -f "src/main/resources/disease_list.json" ]; then
    echo "Creating empty disease_list.json..."
    echo "[]" > src/main/resources/disease_list.json
fi

# Check if .env file exists
if [ ! -f ".env" ]; then
    echo ".env file is missing. Creating from .env.example..."
    if [ -f ".env.example" ]; then
        cp .env.example .env
        echo "   Created .env from .env.example. Please edit it with your API key."
    else
        echo "WARNING: .env.example not found. Creating minimal .env file..."
        echo "ANTHROPIC_API_KEY=your_api_key_here" > .env
    fi
fi

# Stop any running containers
echo "Stopping any running containers..."
docker-compose down || true

# Rebuild the containers
echo "Building Docker containers..."
docker-compose build

# Start the Synthea container
echo "Starting Synthea container..."
docker-compose up -d synthea

# Give the container a moment to start
echo "Waiting for Synthea container to be ready..."
sleep 10

# Check if the container is healthy
CONTAINER_ID=$(docker-compose ps -q synthea)
if [ -z "$CONTAINER_ID" ]; then
    echo "ERROR: Failed to start Synthea container."
    exit 1
fi

# Check if the container is using proper health checks
HEALTH_STATUS=$(docker inspect --format='{{.State.Health.Status}}' $CONTAINER_ID 2>/dev/null || echo "unknown")
if [ "$HEALTH_STATUS" != "healthy" ] && [ "$HEALTH_STATUS" != "unknown" ]; then
    echo "WARNING: Synthea container is not reporting as healthy (status: $HEALTH_STATUS)."
    echo "It may take some time for the container to become fully operational."
    echo "You can proceed, but the pipeline may fail if the container is not ready."
fi

# Check if we should generate modules
if [ "$1" == "--generate-modules" ]; then
    echo "Starting module generator..."
    docker-compose --profile generator up module-generator
fi

echo "=================================================="
echo "Environment is ready!"
echo ""
echo "To generate patients for a disease module:"
echo "nextflow run main.nf --disease_name \"Disease Name\" --generate_patients true"
echo ""
echo "To generate a new module and patients:"
echo "nextflow run main.nf --disease_name \"Disease Name\" --force_generate true --generate_patients true"
echo "=================================================="