#!/bin/bash
set -e

# Check and create .env file if it doesn't exist
if [ ! -f /app/data/.env ]; then
    echo "Creating default .env file..."
    cp /app/.env.template /app/data/.env
    chown elasticsearch:elasticsearch /app/data/.env
    chmod 600 /app/data/.env
fi

# Load environment variables from .env file if it exists
if [ -f /app/data/.env ]; then
    echo "Loading environment variables from .env file..."
    set -o allexport
    source /app/data/.env
    set +o allexport
else
    echo "No .env file found, using default environment variables"
fi

# Set constants
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}

# Set proper permissions early
echo "Setting up directory permissions..."
chown -R elasticsearch:elasticsearch /app/data

# Handle password management
setup_password() {
    # Check if password already exists
    if [ -f /app/data/secrets/elastic_password ]; then
        # Use -r flag to prevent backslash interpretation
        ELASTIC_PASSWORD=$(cat /app/data/secrets/elastic_password)
        echo "Using existing Elasticsearch password."
    else
        # Generate a more container-safe password (alphanumeric only)
        # Avoid special characters that could cause issues with command interpretation
        ELASTIC_PASSWORD=$(tr -dc 'A-Za-z0-9' < /dev/urandom | head -c 20)
        echo "Generated new secure password for Elasticsearch."
        
        # Store password
        echo "$ELASTIC_PASSWORD" > /app/data/secrets/elastic_password
        chmod 600 /app/data/secrets/elastic_password
        chown elasticsearch:elasticsearch /app/data/secrets/elastic_password
    fi
    
    # Make password available to other functions
    export 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
    
    # Ensure Java symlinks have correct ownership
    chown -R elasticsearch:elasticsearch /app/data/jdk
}

# Configure Elasticsearch
configure_elasticsearch() {
    cd $ES_HOME
    
    # Double-check permissions on config directory before proceeding
    echo "Verifying config directory permissions..."
    ls -la $ES_PATH_CONF
    chmod 755 $ES_PATH_CONF
    chown -R elasticsearch:elasticsearch $ES_PATH_CONF
    
    # Add bootstrap password to keystore
    echo "Setting bootstrap password..."
    
    # Verify keystore exists before trying to add password
    if [ ! -f $ES_PATH_CONF/elasticsearch.keystore ]; then
        echo "ERROR: Keystore not found, cannot add bootstrap password. Creating keystore first..."
        setup_keystore || {
            echo "CRITICAL ERROR: Failed to create keystore, cannot proceed."
            exit 1
        }
    fi
    
    # Now add the bootstrap password
    if ! printf "%s" "$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; then
        echo "ERROR: Failed to add bootstrap password to keystore."
        exit 1
    fi
    
    # 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
        chown -R elasticsearch:elasticsearch $ES_PATH_CONF
    fi
    
    # CRITICAL FIX: Remove any index-level settings from elasticsearch.yml to prevent startup failure
    if [ -f $ES_PATH_CONF/elasticsearch.yml ]; then
        echo "Checking elasticsearch.yml for index-level settings..."
        # Create a temporary file
        TEMP_FILE=$(mktemp)
        
        # Filter out any index.* settings
        grep -v "^index\." $ES_PATH_CONF/elasticsearch.yml > $TEMP_FILE
        
        # Also remove other known problematic settings for 8.x
        grep -v "^processors:" $TEMP_FILE > $TEMP_FILE.2 && mv $TEMP_FILE.2 $TEMP_FILE
        grep -v "^bootstrap.system_call_filter:" $TEMP_FILE > $TEMP_FILE.2 && mv $TEMP_FILE.2 $TEMP_FILE
        
        # Add warning comment
        echo "" >> $TEMP_FILE
        echo "# NOTE: The following settings have been removed from this file:" >> $TEMP_FILE
        echo "# - All index.* settings (applied via index templates)" >> $TEMP_FILE
        echo "# - processors setting (no longer supported in 8.x)" >> $TEMP_FILE
        echo "# - bootstrap.system_call_filter (no longer supported in 8.x)" >> $TEMP_FILE
        echo "# See the create_index_template function in start.sh for details" >> $TEMP_FILE
        
        # Replace the original file
        cat $TEMP_FILE > $ES_PATH_CONF/elasticsearch.yml
        rm $TEMP_FILE
        
        # Ensure proper ownership
        chown elasticsearch:elasticsearch $ES_PATH_CONF/elasticsearch.yml
        echo "Cleaned elasticsearch.yml configuration file"
    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
        chown elasticsearch:elasticsearch $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
        chown elasticsearch:elasticsearch /tmp/elastic-certs
        
        ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-certutil ca \
            --out /tmp/elastic-certs/elastic-stack-ca.p12 \
            --pass "cloudron" \
            --silent
            
        ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-certutil cert \
            --ca /tmp/elastic-certs/elastic-stack-ca.p12 \
            --ca-pass "cloudron" \
            --out $ES_PATH_CONF/elastic-certificates.p12 \
            --pass "cloudron" \
            --silent
            
        chown elasticsearch:elasticsearch $ES_PATH_CONF/elastic-certificates.p12
        chmod 600 $ES_PATH_CONF/elastic-certificates.p12
        
        # Make sure we update the keystore with the correct password after generating certificates
        echo "Updating keystore with the new certificate password..."
        setup_keystore
    fi
    
    # Create users file if needed
    if [ ! -f $ES_PATH_CONF/users ]; then
        echo "Creating users file..."
        # Note: We'll reset the password after Elasticsearch starts, so this default password
        # is just for initial bootstrap
        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
    
    # Final permission check
    echo "Final permission check on all data directories..."
    chown -R elasticsearch:elasticsearch /app/data
    chmod 755 /app/data /app/data/config
}

# Create index template with optimized settings
create_index_template() {
    echo "Creating default index template with optimized settings..."
    
    # Wait a moment to ensure Elasticsearch is fully operational
    sleep 10
    
    # Define the template JSON with all the index settings that were previously in elasticsearch.yml
    template_json=$(cat <<EOF
{
  "index_patterns": ["*"],
  "template": {
    "settings": {
      "number_of_shards": 1,
      "number_of_replicas": 0,
      "merge": {
        "scheduler": {
          "max_thread_count": 1
        },
        "policy": {
          "floor_segment": "2mb",
          "max_merge_at_once": 4,
          "segments_per_tier": 8
        }
      }
    }
  },
  "priority": 1,
  "_meta": {
    "description": "Default template with Cloudron optimized settings"
  }
}
EOF
)
    
    # Try multiple times in case Elasticsearch is still initializing
    for i in {1..5}; do
        echo "Attempt $i to create index template..."
        
        # Apply the template
        response=$(curl -s -w "\n%{http_code}" -X PUT "http://localhost:9200/_index_template/cloudron_defaults" \
          -H "Content-Type: application/json" \
          -u "elastic:$ELASTIC_PASSWORD" \
          -d "$template_json")
          
        http_code=$(echo "$response" | tail -n1)
        response_body=$(echo "$response" | sed '$d')
        
        if [ "$http_code" = "200" ] || [ "$http_code" = "201" ]; then
            echo "✅ Index template created successfully with HTTP code $http_code"
            echo "Response: $response_body"
            return 0
        else
            echo "⚠️ Failed to create index template on attempt $i. HTTP code: $http_code"
            echo "Response: $response_body"
            sleep 5
        fi
    done
    
    echo "⚠️ Warning: Failed to create index template after multiple attempts."
    echo "Default settings may not be applied to new indices, but Elasticsearch will still function."
    echo "You can manually create the template later using Elasticsearch API."
    
    # Don't fail the startup process if template creation fails
    return 0
}

# Set system limits - be more tolerant of container restrictions
set_system_limits() {
    echo "Setting system limits for Elasticsearch..."
    
    # Try to set file descriptor limit, but don't fail if it doesn't work
    ulimit -n 65536 2>/dev/null || echo "Warning: Could not set file descriptor limit (not critical)"
    
    # Try to set memory lock limit, but don't fail if it doesn't work
    ulimit -l unlimited 2>/dev/null || echo "Warning: Could not set memory lock limit (not critical)"
    
    # Only try to update transparent huge pages if the file exists and is writable
    if [ -w /sys/kernel/mm/transparent_hugepage/enabled ]; then
        echo never > /sys/kernel/mm/transparent_hugepage/enabled 2>/dev/null || true
    else
        echo "Warning: Cannot modify transparent hugepage settings (read-only filesystem, not critical)"
    fi
    
    # Only try to update vm.max_map_count if sysctl is available and we have permission
    if command -v sysctl >/dev/null && [ $(id -u) -eq 0 ]; then
        sysctl -w vm.max_map_count=262144 2>/dev/null || echo "Warning: Could not set vm.max_map_count (not critical)"
    else
        echo "Warning: Could not set vm.max_map_count (not running as root or sysctl not available)"
    fi
    
    # Add a note about bootstrap.memory_lock if we couldn't set the memory lock
    if ! ulimit -l unlimited 2>/dev/null; then
        echo "Note: Memory locking unavailable. Setting bootstrap.memory_lock=false in elasticsearch.yml"
        if grep -q "bootstrap.memory_lock:" $ES_PATH_CONF/elasticsearch.yml; then
            sed -i 's/bootstrap.memory_lock:.*/bootstrap.memory_lock: false/' $ES_PATH_CONF/elasticsearch.yml
        else
            echo "bootstrap.memory_lock: false" >> $ES_PATH_CONF/elasticsearch.yml
        fi
    fi
}

# Add secure settings to the keystore
setup_keystore() {
    echo "Setting up Elasticsearch keystore with secure settings..."
    
    # Create or recreate the keystore if needed
    if [ ! -f $ES_PATH_CONF/elasticsearch.keystore ]; then
        echo "Creating new Elasticsearch keystore..."
        su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore create" elasticsearch
        
        # Verify keystore was created
        if [ ! -f $ES_PATH_CONF/elasticsearch.keystore ]; then
            echo "ERROR: Failed to create keystore!"
            return 1
        fi
    fi
    
    # Add the certificate passwords to the keystore (as the elasticsearch user)
    echo "Adding certificate passwords to keystore..."
    echo "cloudron" | su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore add -f xpack.security.transport.ssl.keystore.secure_password --stdin" elasticsearch || {
        echo "ERROR: Failed to add keystore password to keystore. Will try to recreate 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
        echo "cloudron" | su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore add -f xpack.security.transport.ssl.keystore.secure_password --stdin" elasticsearch || {
            echo "CRITICAL ERROR: Could not add keystore password to keystore after recreation."
            return 1
        }
    }
    
    echo "cloudron" | su -c "ES_PATH_CONF=$ES_PATH_CONF ES_JAVA_HOME=/app/data/jdk $ES_HOME/bin/elasticsearch-keystore add -f xpack.security.transport.ssl.truststore.secure_password --stdin" elasticsearch || {
        echo "ERROR: Failed to add truststore password to keystore."
        return 1
    }
    
    # Verify keystore permissions
    chmod 600 $ES_PATH_CONF/elasticsearch.keystore
    chown elasticsearch:elasticsearch $ES_PATH_CONF/elasticsearch.keystore
    
    echo "✅ Elasticsearch keystore updated with secure settings"
    return 0
}

# 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"
    # Add explicit settings for transport SSL
    ES_START_CMD="$ES_START_CMD -E xpack.security.transport.ssl.enabled=true"
    ES_START_CMD="$ES_START_CMD -E xpack.security.transport.ssl.verification_mode=certificate"
    ES_START_CMD="$ES_START_CMD -E xpack.security.transport.ssl.keystore.path=elastic-certificates.p12"
    ES_START_CMD="$ES_START_CMD -E xpack.security.transport.ssl.truststore.path=elastic-certificates.p12"
    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 -u "elastic:$ELASTIC_PASSWORD" 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 --password "$ELASTIC_PASSWORD" --url "http://localhost:9200" || true
    
    # Create index template with the settings we removed from elasticsearch.yml
    echo "Applying index templates with settings removed from elasticsearch.yml..."
    create_index_template
    
    # Create credentials file
    cat > /app/data/credentials.txt << EOL
Elasticsearch credentials:
URL: http://localhost:9200
User: elastic
Password: $ELASTIC_PASSWORD
EOL
    chmod 600 /app/data/credentials.txt
    
    echo "-----------------------------"
    echo "Elasticsearch is ready to use!"
    echo "URL: http://localhost:9200"
    echo "User: elastic"
    echo "Password: $ELASTIC_PASSWORD"
    echo "Password is stored in: /app/data/credentials.txt"
    echo "-----------------------------"
}

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

# Final permission check
echo "Performing final permission check on all directories..."
chown -R elasticsearch:elasticsearch /app/data/{elasticsearch,logs,config,jdk,run,secrets}
chmod -R 755 /app/data
chmod 600 /app/data/secrets/elastic_password

set_system_limits
configure_heap
start_elasticsearch

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