#!/bin/bash
set -e

# Change directory to script's location
cd "$(dirname "$0")"

DROP_DB=false
while [[ "$#" -gt 0 ]]; do
    case $1 in
        --drop) DROP_DB=true ;;
    esac
    shift
done

# Source credentials from db_credentials file
if [ -f "./db_credentials" ]; then
    echo "Loading credentials from db_credentials..."
    source ./db_credentials
else
    echo "Error: db_credentials file not found in $(pwd)"
    exit 1
fi

# 1. Provide the PEM key for AWS RDS if not exists
if [ -n "$DB_PEM_URL" ] && [ ! -f "$DB_SSL_ROOT_CERT" ]; then
    echo "Downloading SSL root certificate from AWS..."
    wget -qO "$DB_SSL_ROOT_CERT" "$DB_PEM_URL"
fi

# Ensure PEM has right permissions if it exists
if [ -f "$DB_SSL_ROOT_CERT" ]; then
    chmod 600 "$DB_SSL_ROOT_CERT"
fi

# 2. Login as the RDS admin user (postgres) to create the zabbix user/database
echo "Connecting to PostgreSQL to create Zabbix user and database..."

export PGPASSWORD="$DB_PASSWORD"

# Create the zabbix user if it doesn't already exist
psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -v ON_ERROR_STOP=1 <<EOF
DO \$\$
BEGIN
    IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '$ZBX_DB_USER') THEN
        CREATE ROLE $ZBX_DB_USER WITH LOGIN PASSWORD '$ZBX_DB_PASSWORD';
    END IF;
END
\$\$;
EOF

echo "User '$ZBX_DB_USER' verified/created."

# Create the zabbix database if it doesn't already exist
DB_EXISTS=$(psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -t -c "SELECT 1 FROM pg_database WHERE datname='$ZBX_DB_NAME'" | tr -d '[:space:]')

if [ "$DROP_DB" = "true" ] && [ "$DB_EXISTS" = "1" ]; then
    echo -e "\n========================================"
    echo -e "              WARNING!                  "
    echo -e "========================================"
    echo -e "You requested to completely DROP and RE-INITIATE the database '$ZBX_DB_NAME'."
    echo -e "This will delete ALL data. Are you sure you want to proceed?"
    read -p "Type 'yes' to proceed: " confirm_drop
    if [ "$confirm_drop" != "yes" ]; then
        echo "Database drop cancelled. Exiting."
        exit 1
    fi
    echo "Terminating active connections and dropping database..."
    psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '$ZBX_DB_NAME' AND pid <> pg_backend_pid();"
    psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -c "DROP DATABASE $ZBX_DB_NAME;"
    DB_EXISTS=""
fi

if [ "$DB_EXISTS" != "1" ]; then
    echo "Database '$ZBX_DB_NAME' does not exist. Creating..."
    psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -c "CREATE DATABASE $ZBX_DB_NAME OWNER $ZBX_DB_USER;"
else
    echo "Database '$ZBX_DB_NAME' already exists."
fi

# Grant necessary permissions
psql "host=$DB_HOST port=$DB_PORT dbname=$DB_NAME user=$DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT" -c "GRANT ALL PRIVILEGES ON DATABASE $ZBX_DB_NAME TO $ZBX_DB_USER;"

echo ""
echo "================================================================================"
echo "✅ Initialization Successful!"
echo "================================================================================"
echo "You can now use these settings in your Zabbix server configuration:"
echo "--------------------------------------------------------------------------------"
echo "DBHost=$DB_HOST"
echo "DBName=$ZBX_DB_NAME"
echo "DBUser=$ZBX_DB_USER"
echo "DBPassword=$ZBX_DB_PASSWORD"
echo "DBPort=$DB_PORT"
echo "DBTLSConnect=verify_full"
echo "DBTLSCAFile=$(realpath $DB_SSL_ROOT_CERT)"
echo "================================================================================"
echo ""
echo "To connect manually for testing directly to the Zabbix DB:"
echo "export PGPASSWORD=\"$ZBX_DB_PASSWORD\""
echo "psql \"host=$DB_HOST port=$DB_PORT dbname=$ZBX_DB_NAME user=$ZBX_DB_USER sslmode=$DB_SSL_MODE sslrootcert=$DB_SSL_ROOT_CERT\""
echo ""