#!/bin/bash

# Function to read configuration using bashio if available
read_config() {
    if command -v bashio >/dev/null 2>&1; then
        bashio::config "$1"
    else
        # Fallback for testing outside Home Assistant
        case "$1" in
            "node_role") echo "master" ;;
            "cluster_name") echo "hass-opensearch-cluster" ;;
            "node_name") echo "os-node-1" ;;
            "heap_size") echo "1g" ;;
            "http_port") echo "9200" ;;
            "transport_port") echo "9300" ;;
            "discovery_seed_hosts") echo '["os-master:9300"]' ;;
            "initial_master_nodes") echo '["os-master"]' ;;
            "security_enabled") echo "true" ;;
            "admin_password") echo "SecurePass123!" ;;
            "plugins") echo "[]" ;;
            *) echo "" ;;
        esac
    fi
}

# Read configuration values
NODE_ROLE=$(read_config 'node_role' | tr -d '[]"' | tr ',' ' ')
CLUSTER_NAME=$(read_config 'cluster_name')
NODE_NAME=$(read_config 'node_name')
HEAP_SIZE=$(read_config 'heap_size')
HTTP_PORT=$(read_config 'http_port')
TRANSPORT_PORT=$(read_config 'transport_port')
DISCOVERY_SEEDS=$(read_config 'discovery_seed_hosts' | tr -d '[]"' | tr ',' ' ')
INITIAL_MASTERS=$(read_config 'initial_master_nodes' | tr -d '[]"' | tr ',' ' ')
SECURITY_ENABLED=$(read_config 'security_enabled')
ADMIN_PASSWORD=$(read_config 'admin_password')

# Export environment variables
export OPENSEARCH_INITIAL_ADMIN_PASSWORD="${ADMIN_PASSWORD}"
export DISCOVERY_TYPE="multi-node"

# Set heap size using environment variable
CLEAN_HEAP_SIZE=$(echo "$HEAP_SIZE" | tr -d ' ' | grep -oE '[0-9]+[mg]')
if [ -z "$CLEAN_HEAP_SIZE" ]; then
    CLEAN_HEAP_SIZE="2g"
fi
export OPENSEARCH_JAVA_OPTS="-Xms${CLEAN_HEAP_SIZE} -Xmx${CLEAN_HEAP_SIZE}"

# Create jvm.options directory if it doesn't exist
mkdir -p /usr/share/opensearch/config/jvm.options.d

# Create or restore the main jvm.options file if it's missing
if [ ! -f "/usr/share/opensearch/config/jvm.options" ]; then
    # Create a basic jvm.options file
    cat > /usr/share/opensearch/config/jvm.options << 'EOF'
## JVM configuration

################################################################
## IMPORTANT: JVM heap size
################################################################
##
## You should always set the min and max JVM heap
## size to the same value. For example, to set
## the heap to 4 GB, set:
##
## -Xms4g
## -Xmx4g
##
## See https://opensearch.org/docs/latest/opensearch/install/important-settings/
## for more information
##
################################################################

# Xms represents the initial size of total heap space
# Xmx represents the maximum size of total heap space

-Xms2g
-Xmx2g

################################################################
## Expert settings
################################################################
##
## All settings below this section are considered
## expert settings. Don't tamper with them unless
## you understand what you are doing
##
################################################################

## GC configuration
-XX:+UseConcMarkSweepGC
-XX:CMSInitiatingOccupancyFraction=75
-XX:+UseCMSInitiatingOccupancyOnly

## optimizations

# disable calls to System#gc
-XX:+DisableExplicitGC

# pre-touch memory pages used by the JVM during initialization
-XX:+AlwaysPreTouch

## basic

# force the server VM (remove on 32-bit client JVMs)
-server

# explicitly set the stack size
-Xss1m

# set to headless, just in case
-Djava.awt.headless=true

# ensure UTF-8 encoding by default (e.g. filenames)
-Dfile.encoding=UTF-8

# use our provided JNA always versus the system one
-Djna.nosys=true

# turn off a JDK optimization that throws away stack traces for common
# exceptions because stack traces are important for debugging
-XX:-OmitStackTraceInFastThrow

# flags to configure Netty
-Dio.netty.noUnsafe=true
-Dio.netty.noKeySetOptimization=true
-Dio.netty.recycler.maxCapacityPerThread=0

# log4j 2
-Dlog4j.shutdownHookEnabled=false
-Dlog4j2.disable.jmx=true

# enable JDK 9+ usage and allow the use of the bundle JDK
--add-opens=java.base/java.io=ALL-UNNAMED
-Des.netty.tryReflectionSetAccessible=true

-Djava.locale.providers=COMPAT

# enable the module path
--module-path=${OPENSEARCH_HOME}/lib/*

# allow access to the jdk.unsupported module for sun.misc.Unsafe
--add-modules=jdk.unsupported

# allow access to the jdk.incubator.vector module for VectorUtil
--add-modules=jdk.incubator.vector
EOF
fi

# Export variables for envsubst
export CLUSTER_NAME="${CLUSTER_NAME:-hass-opensearch-cluster}"
export NODE_NAME="${NODE_NAME:-os-node-1}"
export NODE_ROLE="${NODE_ROLE:-master}"
export HTTP_PORT="${HTTP_PORT:-9200}"
export TRANSPORT_PORT="${TRANSPORT_PORT:-9300}"
export DISCOVERY_SEEDS="${DISCOVERY_SEEDS:-os-master:9300}"
export INITIAL_MASTERS="${INITIAL_MASTERS:-os-master}"
export SECURITY_ENABLED="${SECURITY_ENABLED:-false}"

# Generate opensearch.yml from template
envsubst < /etc/opensearch/opensearch.yml.template > /usr/share/opensearch/config/opensearch.yml

# Install plugins if specified
PLUGINS_JSON=$(read_config 'plugins')
if [ -n "$PLUGINS_JSON" ] && [ "$PLUGINS_JSON" != "[]" ] && command -v bashio >/dev/null 2>&1; then
    PLUGINS=$(echo "$PLUGINS_JSON" | jq -r '.[]' 2>/dev/null || echo "")
    if [ -n "$PLUGINS" ]; then
        for plugin in $PLUGINS; do
            echo "Installing plugin: $plugin"
            /usr/share/opensearch/bin/opensearch-plugin install --batch "$plugin" || true
        done
    fi
fi

# Try to set vm.max_map_count if we have privileges
if [ -w /proc/sys/vm/max_map_count ]; then
    echo "Setting vm.max_map_count to 262144"
    echo 262144 > /proc/sys/vm/max_map_count
elif command -v sysctl >/dev/null 2>&1; then
    echo "Trying to set vm.max_map_count with sysctl"
    sysctl -w vm.max_map_count=262144 || true
fi

# Check current value
if [ -r /proc/sys/vm/max_map_count ]; then
    CURRENT_MAP_COUNT=$(cat /proc/sys/vm/max_map_count)
    echo "Current vm.max_map_count: $CURRENT_MAP_COUNT"
    if [ "$CURRENT_MAP_COUNT" -lt 262144 ]; then
        echo "WARNING: vm.max_map_count is too low ($CURRENT_MAP_COUNT), should be at least 262144"
        echo "You may need to run: sudo sysctl -w vm.max_map_count=262144 on the host system"
    fi
fi

# Start OpenSearch
exec /usr/share/opensearch/bin/opensearch