#!/bin/bash
set -e

# Source environment variables
source /app/.env

# Set constants
ELASTIC_PASSWORD="cloudron"
ES_HOME=/usr/share/elasticsearch
ES_PATH_CONF=/app/data/config
export ES_HOME ES_PATH_CONF

# Create directory structure
mkdir -p /app/data/{elasticsearch,logs/gc,config,run,secrets,jdk/bin}
echo "$ELASTIC_PASSWORD" > /app/data/secrets/elastic_password

# Set up Java environment
setup_java() {
    if [ -L /app/data/jdk/bin/java ]; then
        return 0
    fi
    
    echo "Setting up Java environment..."
    
    # Try Java 17 first, then fall back to system Java
    if [ -f /usr/lib/jvm/java-17-openjdk-amd64/bin/java ]; then
        ln -sf /usr/lib/jvm/java-17-openjdk-amd64/bin/{java,javac,javadoc,jar} /app/data/jdk/bin/
    else
        JAVA_PATH=$(which java)
        if [ -n "$JAVA_PATH" ]; then
            ln -sf $JAVA_PATH /app/data/jdk/bin/java
            for tool in javac javadoc jar; do
                ln -sf $(which $tool 2>/dev/null) /app/data/jdk/bin/$tool 2>/dev/null || true
            done
        else
            echo "ERROR: No Java found on system. Elasticsearch requires Java to run."
            exit 1
        fi
    fi
    
    # Verify Java is available
    if [ -L /app/data/jdk/bin/java ]; then
        echo "Java version: $(/app/data/jdk/bin/java -version 2>&1 | head -n 1)"
    else
        echo "ERROR: Failed to link Java executable"
        exit 1
    fi
}

# Configure Elasticsearch
configure_elasticsearch() {
    cd $ES_HOME
    
    # Handle keystore creation and password
    if [ ! -f $ES_PATH_CONF/elasticsearch.keystore ] || [ "$1" = "force" ]; then
        echo "Creating Elasticsearch keystore..."
        [ -f $ES_PATH_CONF/elasticsearch.keystore ] && rm -f $ES_PATH_CONF/elasticsearch.keystore
        su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore create" elasticsearch
    fi
    
    # Add bootstrap password to keystore
    echo "Setting bootstrap password..."
    echo "$ELASTIC_PASSWORD" | su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore add -f -x 'bootstrap.password' --stdin" elasticsearch
    
    # Copy configuration files if needed
    if [ ! -f $ES_PATH_CONF/elasticsearch.yml ]; then
        echo "Setting up configuration files..."
        cp -r $ES_HOME/config/* $ES_PATH_CONF/ || true
        cp /app/elasticsearch.yml $ES_PATH_CONF/elasticsearch.yml || true
    fi
    
    # Update JVM options for GC logs
    if [ -f $ES_PATH_CONF/jvm.options ]; then
        echo "Updating JVM options..."
        sed -i 's|logs/gc.log|/app/data/logs/gc/gc.log|g' $ES_PATH_CONF/jvm.options
    fi
    
    # Generate SSL certificates if needed
    if [ ! -f $ES_PATH_CONF/elastic-certificates.p12 ]; then
        echo "Generating self-signed certificates..."
        
        mkdir -p /tmp/elastic-certs
        ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-certutil ca \
            --out /tmp/elastic-certs/elastic-stack-ca.p12 \
            --pass "" \
            --silent
            
        ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-certutil cert \
            --ca /tmp/elastic-certs/elastic-stack-ca.p12 \
            --ca-pass "" \
            --out $ES_PATH_CONF/elastic-certificates.p12 \
            --pass "" \
            --silent
            
        chown elasticsearch:elasticsearch $ES_PATH_CONF/elastic-certificates.p12
        chmod 600 $ES_PATH_CONF/elastic-certificates.p12
    fi
    
    # Create users file if needed
    if [ ! -f $ES_PATH_CONF/users ]; then
        echo "Creating users file..."
        echo 'elastic:$2a$10$BtVRGAoL8AbgEKnlvYj8cewQF3QkUz1pyL.Ga3j.jFKNUk2yh7.zW' > $ES_PATH_CONF/users
        echo 'kibana_system:$2a$10$BtVRGAoL8AbgEKnlvYj8cewQF3QkUz1pyL.Ga3j.jFKNUk2yh7.zW' >> $ES_PATH_CONF/users
        
        echo 'superuser:elastic' > $ES_PATH_CONF/users_roles
        
        chown elasticsearch:elasticsearch $ES_PATH_CONF/{users,users_roles}
        chmod 600 $ES_PATH_CONF/{users,users_roles}
    fi
    
    # Ensure basic settings in elasticsearch.yml
    for setting in "xpack.security.http.ssl.enabled: false" "network.host: 0.0.0.0" "discovery.type: single-node"; do
        if ! grep -q "$setting" $ES_PATH_CONF/elasticsearch.yml; then
            echo "$setting" >> $ES_PATH_CONF/elasticsearch.yml
        fi
    done
}

# Set system limits
set_system_limits() {
    echo "Setting system limits for Elasticsearch..."
    ulimit -n 65536 || echo "Warning: Could not set file descriptor limit"
    ulimit -l unlimited || echo "Warning: Could not set memory lock limit"
    echo never > /sys/kernel/mm/transparent_hugepage/enabled 2>/dev/null || true
    sysctl -w vm.max_map_count=262144 2>/dev/null || echo "Warning: Could not set vm.max_map_count"
}

# Configure JVM heap size
configure_heap() {
    # Calculate optimal heap size (50% of available memory)
    CONTAINER_MEM=$(cat /sys/fs/cgroup/memory.max 2>/dev/null || echo "4294967296")
    [ "$CONTAINER_MEM" = "max" ] && CONTAINER_MEM="4294967296"  # Default to 4GB if unlimited
    
    HEAP_SIZE=$(expr $CONTAINER_MEM / 2097152)  # Convert to MB and take 50%
    [ $HEAP_SIZE -gt 31744 ] && HEAP_SIZE=31744  # Max 31GB
    [ $HEAP_SIZE -lt 512 ] && HEAP_SIZE=512      # Min 512MB
    
    # Set JVM options
    export ES_JAVA_OPTS="-Xms${HEAP_SIZE}m -Xmx${HEAP_SIZE}m -XX:+UseG1GC -XX:G1ReservePercent=25 -XX:InitiatingHeapOccupancyPercent=30"
    export ES_JAVA_OPTS="$ES_JAVA_OPTS -Xlog:gc*,gc+age=trace,safepoint:file=/app/data/logs/gc/gc.log:utctime,level,pid,tags:filecount=32,filesize=64m"
    
    export PATH=$ES_HOME/bin:$PATH
}

# Start Elasticsearch
start_elasticsearch() {
    # Create PID file
    touch /app/data/run/elasticsearch.pid
    chown elasticsearch:elasticsearch /app/data/run/elasticsearch.pid
    
    # Command to start Elasticsearch
    ES_START_CMD="ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch"
    ES_START_CMD="$ES_START_CMD -E xpack.security.enabled=true -E bootstrap.password=$ELASTIC_PASSWORD"
    ES_START_CMD="$ES_START_CMD -d -p /app/data/run/elasticsearch.pid"
    
    echo "Starting Elasticsearch..."
    cd $ES_HOME
    su -c "$ES_START_CMD" elasticsearch
    
    # Wait for Elasticsearch to start
    echo "Waiting for Elasticsearch to start..."
    attempts=0
    max_attempts=60
    until $(curl --output /dev/null --silent --head --fail http://localhost:9200); do
        if ! ps -p $(cat /app/data/run/elasticsearch.pid 2>/dev/null) > /dev/null 2>&1; then
            echo "ERROR: Elasticsearch process is not running. Logs:"
            cat /app/data/logs/*.log
            exit 1
        fi
        
        printf '.'
        sleep 5
        
        attempts=$((attempts+1))
        if [ $attempts -ge $max_attempts ]; then
            echo "ERROR: Elasticsearch failed to start after 5 minutes. Logs:"
            cat /app/data/logs/*.log
            exit 1
        fi
    done
    echo "Elasticsearch is up and running!"
    
    # Reset the elastic user password
    cd $ES_HOME
    echo "y" | ES_JAVA_HOME=/app/data/jdk bin/elasticsearch-reset-password -u elastic -b -p "$ELASTIC_PASSWORD" --url "http://localhost:9200" || true
    
    # Create credentials file
    cat > /app/data/credentials.txt << EOL
Elasticsearch credentials:
URL: http://localhost:9200
User: elastic
Password: $ELASTIC_PASSWORD
EOL
    
    echo "-----------------------------"
    echo "Elasticsearch is ready to use!"
    echo "URL: http://localhost:9200"
    echo "User: elastic"
    echo "Password: $ELASTIC_PASSWORD"
    echo "-----------------------------"
}

# Main execution flow
setup_java
configure_elasticsearch
[ ! -f /app/data/.initialized ] && touch /app/data/.initialized

# Ensure correct permissions
chown -R elasticsearch:elasticsearch /app/data/{elasticsearch,logs,config,jdk,run,secrets}

set_system_limits
configure_heap
start_elasticsearch

# Keep container running
tail -f /app/data/logs/*.log 2>/dev/null || sleep infinity