#!/bin/bash

VPROTECT_DIR='/opt/vprotect/server'
QUARKUS_PROPERTIES_FILE='quarkus.properties'
MYSQL_DB_NAME='vprotect'
MYSQL_USER_NAME='vprotect'
MYSQL_ROOT_PASSWORD_LENGTH=80
MYSQL_CONFIG_FILE=''
MYSQL_CONFIG_BACKUP_FILE=''
MYSQL_CLIENT=''
MYSQL_ROOT_PASSWORD=''
MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE=''
MYSQL_BASE_MIGRATION_PATH="$VPROTECT_DIR/base_migration.sql"
DB_SERVICE='mariadb'
NON_INTERACTIVE=false
RECREATE_DATABASE=false

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

ECHO () {
echo -e "\e[92m[$1]\e[0m"
}

RECREATE_VPROTECT_DATABASE() {
	if systemctl is-active --quiet vprotect-server; then
		FATAL "The vProtect Server service is active. Stop it before using --recreate-vprotect-database."
	fi

	ECHO "Recreating vProtect database"
	RUN_CRITICAL "Could not drop the vProtect database" MARIADB_AS_ROOT -e "DROP DATABASE IF EXISTS $MYSQL_DB_NAME"
	RUN_CRITICAL "Could not drop existing vProtect database users" MARIADB_AS_ROOT -e "DROP USER IF EXISTS '$MYSQL_USER_NAME'@'localhost', '$MYSQL_USER_NAME'@'%'"
}


FATAL() {
	echo -e "${RED}ERROR:${NC} $*" >&2
	echo "The local MariaDB setup was interrupted. Resolve the reported problem and run the installer again on a clean database." >&2
	exit 1
}

WARN() {
	echo -e "${YELLOW}WARNING:${NC} $*" >&2
}

RUN_CRITICAL() {
	local description="$1"
	shift
	"$@"
	local status=$?
	if [[ $status -ne 0 ]]; then
		FATAL "$description (exit code $status)"
	fi
}

PARSE_ARGUMENTS() {
	while [[ $# -gt 0 ]]; do
		case "$1" in
			--help|-h)
				cat <<EOF_HELP
Usage: $0 [--non-interactive] [--recreate-vprotect-database]

Configures only the local MariaDB instance used by vProtect.

Options:
  --non-interactive
      Read the MariaDB root password from SBR_MYSQL_ROOT_PASSWORD.
  --recreate-vprotect-database
      DESTRUCTIVE: remove the vprotect database and its MariaDB users before
      bootstrapping them again. The vprotect-server service must be stopped.
  -h, --help
      Display this help and exit.
EOF_HELP
				exit 0
				;;
			--non-interactive)
				NON_INTERACTIVE=true
				;;
			--recreate-vprotect-database)
				RECREATE_DATABASE=true
				;;
			*)
				echo "Usage: $0 [--non-interactive] [--recreate-vprotect-database]" >&2
				exit 2
				;;
		esac
		shift
	done
}

REQUIRE_ROOT() {
	if [[ "$EUID" -ne 0 ]]; then
		FATAL "Root permissions are required."
	fi
}

INITIALIZE_MARIADB_CONFIG_FILE() {
	MYSQL_CONFIG_FILE='/etc/my.cnf'
	if [[ "$ID" == "debian" || "${ID_LIKE:-}" == *"debian"* ]]; then
		MYSQL_CONFIG_FILE='/etc/mysql/mariadb.conf.d/99-vprotect.cnf'
	fi

	if [[ ! -f "$MYSQL_CONFIG_FILE" ]] && ! touch "$MYSQL_CONFIG_FILE"; then
		FATAL "Cannot create or modify $MYSQL_CONFIG_FILE."
	fi
}

SELECT_MARIADB_CLIENT() {
	if command -v mariadb &>/dev/null; then
		MYSQL_CLIENT='mariadb'
	elif command -v mysql &>/dev/null; then
		MYSQL_CLIENT='mysql'
	else
		FATAL "Neither mariadb nor mysql client command is available."
	fi
}

MARIADB_AS_ROOT() {
	"$MYSQL_CLIENT" -uroot -p"$MYSQL_ROOT_PASSWORD" "$@"
}

SUDO(){
	ECHO "Configuring sudo"
	if ! grep LIBGUESTFS_BACKEND /etc/sudoers &>/dev/null; then
		if ! sed -Ei 's/(.*env_keep.*LC_ALL)/\1 LIBGUESTFS_BACKEND/' /etc/sudoers; then
			WARN "Could not add LIBGUESTFS_BACKEND to /etc/sudoers. Database setup will continue."
		fi
	fi
}

FILE_ENDS_WITH_NEWLINE() {
	[[ $(tail -c1 "$1" | wc -l) -gt 0 ]]
}

DECODE_PASSWORD() {
	local password="$1"
	local decoded_password

	decoded_password=$(printf '%s' "$password" | base64 -d 2>/dev/null)
	if [[ "${decoded_password:0:8}" == "[BASE64]" ]]; then
		printf '%s' "${decoded_password:8}"
	else
		printf '%s' "$password"
	fi
}

ESCAPE_ROOT_PASSWORD() {
	MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE=$(printf '%s' "$MYSQL_ROOT_PASSWORD" | sed -e 's/\\/\\\\/g' -e "s/'/\\\'/g")
}

ENSURE_PASSWORD() {
	local pass

	stty -echo
	printf "Type password again: "
	IFS= read -rs pass; echo
	pass=$(DECODE_PASSWORD "$pass")
	stty echo

	if [[ "$MYSQL_ROOT_PASSWORD" != "$pass" ]]; then
		echo "Not the same password!"
		return 1
	fi

	ESCAPE_ROOT_PASSWORD
}

SECURE_INSTANCE() {
	if [[ "$NON_INTERACTIVE" != true ]]; then
		ASK_FOR_ROOT_PASSWORD "Set MySQL root password: "
		until [[ $? == 0 ]]
		do
		    ASK_FOR_ROOT_PASSWORD "Set MySQL root password: "
		done
	fi
	ECHO "Removing anonymous users"
	RUN_CRITICAL "Could not remove anonymous MariaDB users" "$MYSQL_CLIENT" -uroot -e "DELETE FROM mysql.user WHERE User='';"
	ECHO "Removing remote root access"
	RUN_CRITICAL "Could not remove remote MariaDB root access" "$MYSQL_CLIENT" -uroot -e "DELETE FROM mysql.user WHERE User='root' AND Host NOT IN ('localhost', '127.0.0.1', '::1');"
	ECHO "Removing test database"
	RUN_CRITICAL "Could not remove the MariaDB test database" "$MYSQL_CLIENT" -uroot -e "DROP DATABASE IF EXISTS test;"
	RUN_CRITICAL "Could not remove test database privileges" "$MYSQL_CLIENT" -uroot -e "DELETE FROM mysql.db WHERE Db='test' OR Db='test\\_%';"
	RUN_CRITICAL "Could not flush MariaDB privileges" "$MYSQL_CLIENT" -uroot -e "FLUSH PRIVILEGES;"
	ECHO "Changing MariaDB root user password"
	RUN_CRITICAL "Could not configure root@localhost" "$MYSQL_CLIENT" -uroot -e "CREATE USER IF NOT EXISTS 'root'@'localhost' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE'; FLUSH PRIVILEGES;"
	RUN_CRITICAL "Could not configure root@127.0.0.1" MARIADB_AS_ROOT -e "CREATE USER IF NOT EXISTS 'root'@'127.0.0.1' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE';"
	RUN_CRITICAL "Could not configure root@::1" MARIADB_AS_ROOT -e "CREATE USER IF NOT EXISTS 'root'@'::1' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE'; FLUSH PRIVILEGES;"
	RUN_CRITICAL "Could not update the password for root@localhost" "$MYSQL_CLIENT" -uroot -e "ALTER USER 'root'@'localhost' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE'; FLUSH PRIVILEGES;"
	RUN_CRITICAL "Could not update the password for root@127.0.0.1" MARIADB_AS_ROOT -e "ALTER USER 'root'@'127.0.0.1' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE';"
	RUN_CRITICAL "Could not update the password for root@::1" MARIADB_AS_ROOT -e "ALTER USER 'root'@'::1' IDENTIFIED BY '$MYSQL_ROOT_PASSWORD_MARIADB_ESCAPE'; FLUSH PRIVILEGES;"

	if grep -q "disable_unix_socket" "$MYSQL_CONFIG_FILE"; then
		ECHO "Unix socket settings OK!"
	else
		RUN_CRITICAL "Could not disable unix socket authentication" DISABLE_UNIX_SOCKET
	fi
}

ENSURE_CONFIG() {
	ECHO "Ensuring MariaDB service is enabled"
	RUN_CRITICAL "Could not enable the MariaDB service" systemctl enable "$DB_SERVICE"
	ECHO "MariaDB service is enabled"
}

set_mysql_param() {
  local file="${MYSQL_CONFIG_FILE:-/etc/mysql/my.cnf}"
  local key value kv

  if [[ $# -eq 1 && "$1" == *"="* ]]; then
    kv="$1"
    key="${kv%%=*}"
    value="${kv#*=}"
  elif [[ $# -ge 2 ]]; then
    key="$1"; value="$2"; kv="$key=$value"
    [[ -n "$3" ]] && file="$3"
  else
    echo "Usage: set_mysql_param KEY=VALUE | set_mysql_param KEY VALUE [FILE]" >&2
    return 2
  fi

  [[ -f "$file" ]] || touch "$file" || return 1
  if [[ -z "$MYSQL_CONFIG_BACKUP_FILE" ]]; then
    MYSQL_CONFIG_BACKUP_FILE="${file}.bak.$(date +%Y%m%d%H%M%S)"
    cp "$file" "$MYSQL_CONFIG_BACKUP_FILE" || return 1
    ECHO "MariaDB configuration backup saved to $MYSQL_CONFIG_BACKUP_FILE"
  fi

  awk -v key="$key" -v value="$value" '
    BEGIN {
      kv = key "=" value
      mysqld = 0
      found = 0
      seen_mysqld = 0
    }
    function print_kv(){ print kv; found=1 }
    {
      if ($0 ~ /^[[:space:]]*\[[^]]+\][[:space:]]*$/) {
        if (mysqld && !found) { print_kv() }
        mysqld = ($0 ~ /^[[:space:]]*\[mysqld\][[:space:]]*$/)
        if (mysqld) seen_mysqld = 1
        print $0
        next
      }
      if (mysqld) {
        if ($0 ~ /^[[:space:]]*[#;]/ || $0 ~ /^[[:space:]]*$/) { print $0; next }
        if ($0 ~ "^[[:space:]]*" key "[[:space:]]*=") {
          if (!found) print kv
          found = 1
          next
        }
      }
      print $0
    }
    END {
      if (!seen_mysqld) {
        print ""
        print "[mysqld]"
      }
      if (!found) print kv
    }
  ' "$file" > "${file}.tmp" || return 1
  mv "${file}.tmp" "$file" || return 1
}

DISABLE_UNIX_SOCKET() {
	ECHO "Disabling unix socket"
	if grep -q "unix_socket=OFF" "$MYSQL_CONFIG_FILE"; then
		sed -i "s/unix_socket=OFF/disable_unix_socket/g" "$MYSQL_CONFIG_FILE" || return 1
	elif grep -q "unix_socket=ON" "$MYSQL_CONFIG_FILE"; then
		sed -i "s/unix_socket=ON/disable_unix_socket/g" "$MYSQL_CONFIG_FILE" || return 1
	elif grep -q "\\[mysqld\\]" "$MYSQL_CONFIG_FILE"; then
		sed -i "s/\\[mysqld\\]/\\[mysqld\\]\\ndisable_unix_socket/g" "$MYSQL_CONFIG_FILE" || return 1
	else
		printf '\n[mysqld]\ndisable_unix_socket\n' >> "$MYSQL_CONFIG_FILE" || return 1
	fi
}

ASK_FOR_ROOT_PASSWORD() {
	stty -echo
	printf "$1"
	IFS= read -rs MYSQL_ROOT_PASSWORD; echo
	MYSQL_ROOT_PASSWORD=$(DECODE_PASSWORD "$MYSQL_ROOT_PASSWORD")
	stty echo

	until [ "$(expr "$MYSQL_ROOT_PASSWORD" : '.*')" -lt "$MYSQL_ROOT_PASSWORD_LENGTH" ]
	do
		echo "Password too long. Max length is 79 characters."
		printf "$1"
		IFS= read -rs MYSQL_ROOT_PASSWORD; echo
		MYSQL_ROOT_PASSWORD=$(DECODE_PASSWORD "$MYSQL_ROOT_PASSWORD")
	done

	ENSURE_PASSWORD
}

CHECK_ROOT_ACCESS() {
	if MARIADB_AS_ROOT -e "exit" &>/dev/null; then
		ECHO "Setting up MariaDB for vProtect"
	else
		FATAL "Cannot authenticate as the MariaDB root user. Verify SBR_MYSQL_ROOT_PASSWORD or the local root authentication configuration."
	fi
}

ENSURE_DATABASE_EXISTS() {
	if MARIADB_AS_ROOT "$MYSQL_DB_NAME" -e "exit" &>/dev/null; then
		ECHO "vProtect database already exists"
	else
		RUN_CRITICAL "Could not create the vProtect database" MARIADB_AS_ROOT -e "CREATE DATABASE $MYSQL_DB_NAME CHARACTER SET utf8 COLLATE utf8_unicode_ci;"
		ECHO "vProtect database created"
	fi
}

CREATE_VPROTECT_USER() {
	MYSQL_USER_PASSWORD=$(< /dev/urandom tr -dc A-Za-z0-9 | head -c32)
	RUN_CRITICAL "Could not create the local vProtect database user. The script requires a fresh installation without an existing vProtect user; if this is a rebuild, stop vprotect-server and rerun with --recreate-vprotect-database" MARIADB_AS_ROOT -e "CREATE USER '$MYSQL_USER_NAME'@'localhost' IDENTIFIED BY '$MYSQL_USER_PASSWORD'"
	RUN_CRITICAL "Could not grant local database privileges to vProtect" MARIADB_AS_ROOT -e "GRANT ALL PRIVILEGES ON $MYSQL_DB_NAME.* TO '$MYSQL_USER_NAME'@'localhost' WITH GRANT OPTION"
	RUN_CRITICAL "Could not create the vProtect database user for TCP connections" MARIADB_AS_ROOT -e "CREATE USER '$MYSQL_USER_NAME'@'%' IDENTIFIED BY '$MYSQL_USER_PASSWORD'"
	RUN_CRITICAL "Could not grant TCP database privileges to vProtect" MARIADB_AS_ROOT -e "GRANT ALL PRIVILEGES ON $MYSQL_DB_NAME.* TO '$MYSQL_USER_NAME'@'%' WITH GRANT OPTION"
	RUN_CRITICAL "Could not flush MariaDB privileges" MARIADB_AS_ROOT -e "FLUSH PRIVILEGES"
	ECHO "vProtect database user created"
	PUT_PASSWORD_IN_QUARKUS_PROPERTIES
}

PUT_PASSWORD_IN_QUARKUS_PROPERTIES() {
	local properties_file="$VPROTECT_DIR/$QUARKUS_PROPERTIES_FILE"

	if [[ ! -f "$properties_file" || ! -w "$properties_file" ]]; then
		FATAL "Cannot update $properties_file with vProtect database credentials."
	fi

	RUN_CRITICAL "Could not remove the previous datasource password from $properties_file" sed -i "/quarkus.datasource.password=/d" "$properties_file"
	RUN_CRITICAL "Could not remove the previous datasource username from $properties_file" sed -i "/quarkus.datasource.username=/d" "$properties_file"
	if ! FILE_ENDS_WITH_NEWLINE "$properties_file"; then
		RUN_CRITICAL "Could not append a newline to $properties_file" sh -c 'printf "\\n" >> "$1"' _ "$properties_file"
	fi
	RUN_CRITICAL "Could not save the vProtect datasource username in $properties_file" sh -c 'printf "%s\\n" "$2" >> "$1"' _ "$properties_file" "quarkus.datasource.username=$MYSQL_USER_NAME"
	RUN_CRITICAL "Could not save the vProtect datasource password in $properties_file" sh -c 'printf "%s\\n" "$2" >> "$1"' _ "$properties_file" "quarkus.datasource.password=$MYSQL_USER_PASSWORD"
	ECHO "Randomly generated password for vProtect database user saved"
}

INSERT_BASE_MIGRATION() {
	RUN_CRITICAL "Could not import the base vProtect schema from $MYSQL_BASE_MIGRATION_PATH" MARIADB_AS_ROOT -e "source $MYSQL_BASE_MIGRATION_PATH" "$MYSQL_DB_NAME"
}

VALIDATE_LOWER_CASE_TABLE_NAMES() {
	local lower_case_table_names

	if ! lower_case_table_names=$(MARIADB_AS_ROOT -s -N -e "SELECT @@GLOBAL.lower_case_table_names;"); then
		FATAL "Cannot read MariaDB lower_case_table_names. Verify root access and the local MariaDB service."
	fi

	if [[ "$lower_case_table_names" != "1" ]]; then
		echo "ERROR: MariaDB lower_case_table_names is '$lower_case_table_names'; vProtect requires '1'." >&2
		echo "ERROR: This value must be set before MariaDB initializes its data directory." >&2
		echo "ERROR: Verify $MYSQL_CONFIG_FILE and any custom hardening policy, then initialize MariaDB with lower_case_table_names=1 before running this script." >&2
		exit 1
	fi

	ECHO "MariaDB lower_case_table_names=1 verified"
}

PREPARE_NON_INTERACTIVE_ROOT_PASSWORD() {
	if [[ -z "$MYSQL_ROOT_PASSWORD" ]]; then
		echo "SBR_MYSQL_ROOT_PASSWORD must be set in non-interactive mode."
		return 1
	fi

	MYSQL_ROOT_PASSWORD=$(DECODE_PASSWORD "$MYSQL_ROOT_PASSWORD")
	if [[ -z "$MYSQL_ROOT_PASSWORD" ]]; then
		echo "SBR_MYSQL_ROOT_PASSWORD must not be empty."
		return 1
	fi

	if [ "$(expr "$MYSQL_ROOT_PASSWORD" : '.*')" -ge "$MYSQL_ROOT_PASSWORD_LENGTH" ]; then
		echo "Password too long. Max length is 79 characters."
		return 1
	fi

	ESCAPE_ROOT_PASSWORD
}

CONFIGURE_MARIADB_TUNING() {
	local parameter parameter_name
	local -a parameters=(
		"innodb_lock_wait_timeout=120"
		"innodb_buffer_pool_size=2G"
		"innodb_log_file_size=512M"
		"max_connections=320"
		"innodb_flush_log_at_trx_commit=2"
		"query_cache_type=0"
		"query_cache_size=0"
		"thread_cache_size=64"
		"join_buffer_size=1M"
		"sort_buffer_size=1M"
		"read_buffer_size=1M"
		"read_rnd_buffer_size=1M"
		"innodb_file_per_table=1"
		"innodb_buffer_pool_instances=2"
		"tmp_table_size=256M"
		"max_heap_table_size=256M"
		"innodb_flush_method=O_DIRECT"
	)

	for parameter in "${parameters[@]}"; do
		parameter_name="${parameter%%=*}"
		RUN_CRITICAL "Could not configure $parameter_name" set_mysql_param "$parameter"
	done
}

main() {
	local script_dir

	script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
	cd "$script_dir" || FATAL "Cannot change directory to $script_dir."

	PARSE_ARGUMENTS "$@"
	REQUIRE_ROOT
	. /etc/os-release
	INITIALIZE_MARIADB_CONFIG_FILE
	SELECT_MARIADB_CLIENT
	MYSQL_ROOT_PASSWORD="${SBR_MYSQL_ROOT_PASSWORD:-}"

	if [[ "$NON_INTERACTIVE" == true ]]; then
		PREPARE_NON_INTERACTIVE_ROOT_PASSWORD || exit 1
	fi

	SUDO
	RUN_CRITICAL "Could not configure lower_case_table_names=1 in $MYSQL_CONFIG_FILE" set_mysql_param "lower_case_table_names=1"
	ENSURE_CONFIG

	ECHO "Restarting MariaDB after applying lower_case_table_names"
	if ! systemctl restart "$DB_SERVICE"; then
		FATAL "MariaDB did not restart after configuring lower_case_table_names. Check systemctl status $DB_SERVICE and the MariaDB error log; a custom hardening policy may override or block this option."
	fi

	if "$MYSQL_CLIENT" -uroot -e "exit" &>/dev/null; then
		SECURE_INSTANCE
	elif [[ "$NON_INTERACTIVE" != true ]]; then
		ASK_FOR_ROOT_PASSWORD "Please provide MySQL root password: "
		until [[ $? == 0 ]]
		do
			ASK_FOR_ROOT_PASSWORD "Please provide MySQL root password: "
		done
	fi

	CHECK_ROOT_ACCESS
	VALIDATE_LOWER_CASE_TABLE_NAMES
	if [[ "$RECREATE_DATABASE" == true ]]; then
		RECREATE_VPROTECT_DATABASE
	fi
	ENSURE_DATABASE_EXISTS
	CREATE_VPROTECT_USER
	CONFIGURE_MARIADB_TUNING

	ECHO "Restarting MariaDB after applying performance settings"
	if ! systemctl restart "$DB_SERVICE"; then
		FATAL "MariaDB did not restart after applying performance settings. Check systemctl status $DB_SERVICE, the MariaDB error log, and $MYSQL_CONFIG_FILE."
	fi

	ECHO "Processing base migration"
	INSERT_BASE_MIGRATION
	ECHO "MariaDB and vProtect database were prepared successfully."
}

main "$@"
