#!/bin/bash
set -eo pipefail

# _log prints a timestamped (ISO 8601), color-coded structured log message.
# The message includes a log level and the message itself.
# If <MESSAGE> is omitted, it reads from stdin, allowing multi-line input.
#
# Arguments:
#   $1 - <LEVEL>   : Log level (e.g., Warn, Error, Debug)
#   $2 - <MESSAGE> : Message to log; if omitted, the function reads from stdin
#
# Usage:
#   _log <LEVEL> [MESSAGE]
#   _log Warn "Disk space low"
#   echo "Database connection lost" | _log Error
#
# Output:
#   2025-10-16T12:34:56+00:00 [Warn] [Entrypoint] Disk space low
#   2025-10-16T12:35:01+00:00 [Error] [Entrypoint] Database connection lost
_log() {
  local level="$1"; shift

  local dt
  dt="$(date --rfc-3339=seconds)"

  local color_reset="\033[0m"
  local color=""
  case "$level" in
  Warn)  color="\033[1;33m" ;; # yellow
  Error) color="\033[1;31m" ;; # red
  Debug) color="\033[1;34m" ;; # blue
  esac

  local msg="$*"
  if [ "$#" -eq 0 ]; then
    msg="$(cat)"
  fi

  printf '%b%s [%s] [Entrypoint] %s%b\n' "$color" "$dt" "$level" "$msg" "$color_reset"
}


# _dbg logs a message of type 'Debug' using log.
_dbg() {
  _log Debug "$@"
}

# mysql_note logs a message of type 'Note' using log.
note() {
  _log Note "$@"
}

# mysql_warn logs a message of type 'Warning' using log and writes to stderr.
warn() {
  _log Warn "$@" >&2
}

# log_error logs a message of type 'ERROR' using log, writes to stderr, prints a container removal hint, and
# exits with status 1.
log_error() {
  _log Error "$@" >&2
  note "Remove this container with 'docker rm -f <container_name>' before retrying"
  exit 1
}

# exec_sql executes a local SQL query, retrying until success or timeout. Ensures reliability during slow
# container or resource startup. On timeout, it prints the provided error prefix followed by filtered Doltgres output.
# Errors are parsed to remove blank lines and extract only relevant error text. Use --show-result to display successful
# query results.
#
# Usage:
#   exec_sql [--show-result] "<ERROR_MESSAGE>" "<QUERY>"
#   exec_sql [--show-result] "<ERROR_MESSAGE>" < /docker-entrypoint-initdb.d/init.sql
#   cat /docker-entrypoint-initdb.d/init.sql | exec_sql [--show-result] "<ERROR_MESSAGE>"
#
# Output:
#   Prints query output only if --show-result is specified.
exec_sql() {
  local show_result=0
  if [ "$1" = "--show-result" ]; then
    show_result=1
    shift
  fi

  local error_message="$1"
  local query="${2:-}"
  local timeout="${DOLTGRES_SERVER_TIMEOUT:-300}"
  local start_time now output status

  start_time=$(date +%s)

  echo "running psql command found in: "
  echo $(which psql)

  local user=$DOLTGRES_USER
  local password=$DOLTGRES_PASSWORD

  # use the default user and pass if not specified in the env
  if [ -z $user ]; then
      user="postgres"
  fi

  if [ -z $password ]; then
      password="password"
  fi
  
  while true; do
      if [ -n "$query" ]; then
          set +e
          output=$(PGPASSWORD=$password psql -h 127.0.0.1 -U $user -c "$query" 2>&1)
          status=$?
          set -e
    else
      set +e # tmp disabled to initdb.d/ file err
      output=$(PGPASSWORD=$password psql -h 127.0.0.1 -U $user < /dev/stdin 2>&1)
      status=$?
      set -e
    fi

    if [ "$status" -eq 0 ]; then
      [ "$show_result" -eq 1 ] && echo "$output" | grep -v "^$" || true
      return 0
    fi

    if echo "$output" | grep -qiE "Error [0-9]+ \([A-Z0-9]+\)"; then
      log_error "$error_message$(echo "$output" | grep -iE "Error|error")"
    fi

    if [ "$timeout" -ne 0 ]; then
      now=$(date +%s)
      if [ $((now - start_time)) -ge "$timeout" ]; then
        log_error "$error_message$(echo "$output" | grep -iE "Error|error" || true)"
      fi
    fi

    sleep 1
  done
}

CONTAINER_DATA_DIR="/var/lib/doltgres"
INIT_COMPLETED="$CONTAINER_DATA_DIR/.init_completed"
SERVER_CONFIG_DIR="/etc/doltgres/servercfg.d"
SERVER_PID=-1

# check_for_doltgres_binary verifies that the dolt binary is present and executable in the system PATH.
# If not found or not executable, it logs an error and exits.
check_for_doltgres_binary() {
  local doltgres_bin
  doltgres_bin=$(which doltgres)
  if [ ! -x "$doltgres_bin" ]; then
    log_error "doltgres binary executable not found"
  fi
}

# get_env_var returns the value of an environment variable, preferring DOLTGRES_* over POSTGRES_*.
# Arguments:
#   $1 - The base variable name (e.g., "USER" for POSTGRES_USER or DOLTGRES_USER)
# Output:
#   Prints the value of the first set variable, or an empty string if neither is set.
get_env_var() {
  local var_name="$1"
  local doltgres_var="DOLTGRES_${var_name}"
  local postgres_var="POSTGRES_${var_name}"

  if [ -n "${!doltgres_var}" ]; then
    echo "${!doltgres_var}"
  elif [ -n "${!postgres_var}" ]; then
    echo "${!postgres_var}"
  else
    echo ""
  fi
}

# get_pgdata_var returns the value of either the DOLTGRES_DATA var or the PGDATA var, in that order.
get_pgdata_var() {
  if [ -n "${DOLTGRES_DATA}" ]; then
    echo "${DOLTGRES_DATA}"
  elif [ -n "${PGDATA}" ]; then
    echo "${PGDATA}"
  else
    echo ""
  fi
}

# get_env_var_name returns the name of the environment variable that is set, preferring DOLTGRES_* over POSTGRES_*.
# Arguments:
#   $1 - The base variable name (e.g., "USER" for POSTGRES_USER or DOLTGRES_USER)
# Output:
#   Prints the name of the first set variable, or both names if neither is set.
get_env_var_name() {
  local var_name="$1"
  local doltgres_var="DOLTGRES_${var_name}"
  local postgres_var="POSTGRES_${var_name}"

  if [ -n "${!doltgres_var}" ]; then
    echo "DOLTGRES_${var_name}"
  elif [ -n "${!postgres_var}" ]; then
    echo "POSTGRES_${var_name}"
  else
    echo "POSTGRES_${var_name}/DOLTGRES_${var_name}"
  fi
}

# get_config_file_path_if_exists checks for config files of a given type in a directory.
# Arguments:
#   $1 - Directory to search in
#   $2 - File type/extension to search for (e.g., 'json', 'yaml')
# Output:
#   Sets CONFIG_PROVIDED to the path of the config file if exactly one is found, or empty otherwise.
#   Logs a warning if multiple config files are found and uses the default config.
get_config_file_path_if_exists() {
  CONFIG_PROVIDED=
  local CONFIG_DIR=$1
  local FILE_TYPE=$2
  if [ -d "$CONFIG_DIR" ]; then
    note "Checking for config provided in $CONFIG_DIR"
    local number_of_files_found
    number_of_files_found=$(find "$CONFIG_DIR" -type f -name "*.$FILE_TYPE" | wc -l)
    if [ "$number_of_files_found" -gt 1 ]; then
      CONFIG_PROVIDED=
      warn "Multiple config files found in $CONFIG_DIR, using default config"
    elif [ "$number_of_files_found" -eq 1 ]; then
      local files_found
      files_found=$(ls "$CONFIG_DIR"/*."$FILE_TYPE")
      note "$files_found file is found"
      CONFIG_PROVIDED=$files_found
    else
      CONFIG_PROVIDED=
    fi
  else
      note "No config dir found in $CONFIG_DIR"
  fi
}

# docker_process_init_files Runs files found in /docker-entrypoint-initdb.d before the server is started.
# Taken from https://github.com/docker-library/mysql/blob/master/8.0/docker-entrypoint.sh
# Usage:
#   docker_process_init_files [file [file ...]]
#   e.g., docker_process_init_files /always-initdb.d/*
# Processes initializer files based on file extensions.
docker_process_init_files() {
  local f
  echo
  for f; do
    case "$f" in
    *.sh)
      if [ -x "$f" ]; then
        note "$0: running $f"
        if ! "$f"; then
          log_error "Failed to execute $f: "
        fi
      else
        note "$0: sourcing $f"
        if ! . "$f"; then
          log_error "Failed to execute $f: "
        fi
      fi
      ;;
    *.sql)
      note "$0: running $f"
      exec_sql --show-result "Failed to execute $f: " < "$f"
      ;;
    *.sql.bz2)
      note "$0: running $f"
      bunzip2 -c "$f" | exec_sql --show-result "Failed to execute $f: "
      ;;
    *.sql.gz)
      note "$0: running $f"
      gunzip -c "$f" | exec_sql --show-result "Failed to execute $f: "
      ;;
    *.sql.xz)
      note "$0: running $f"
      xzcat "$f" | exec_sql --show-result "Failed to execute $f: "
      ;;
    *.sql.zst)
      note "$0: running $f"
      zstd -dc "$f" | exec_sql --show-result "Failed to execute $f: "
      ;;
    *)
      warn "$0: ignoring $f"
      ;;
    esac
    echo
  done
}

# is_port_open checks if a TCP port is open on a given host.
# Arguments:
#   $1 - Host (IP or hostname)
#   $2 - Port number
# Returns:
#   0 if the port is open, non-zero otherwise.
is_port_open() {
  local host="$1"
  local port="$2"
  timeout 1 bash -c "cat < /dev/null > /dev/tcp/$host/$port" &>/dev/null
  return $?
}

# start_server starts the Doltgres server in the background and waits until it is ready to accept connections.
# It manages the server process, restarts it if necessary, and checks for readiness by probing the configured port.
# The function retries until the server is available or a timeout is reached, handling process management and logging.
# Arguments:
#   $@ - Additional arguments to pass to `doltgres`
# Returns:
#   0 if the server starts successfully and is ready to accept connections; exits with error otherwise.
start_server() {
  local timeout="${DOLTGRES_SERVER_TIMEOUT:-300}"
  local start_time
  start_time=$(date +%s)

  # If any of the doltgres or postgres specific user, password, and DB env vars are set, export them
  # to the ones that doltgres understands so they can be used during initialization.
  user=$(get_env_var "USER")
  password=$(get_env_var "PASSWORD")
  db=$(get_env_var "DB")
  export DOLTGRES_USER=$user
  export DOLTGRES_PASSWORD=$password
  
  SERVER_PID=-1

  trap 'note "Caught Ctrl+C, shutting down Doltgres server..."; [ $SERVER_PID -ne -1 ] && kill "$SERVER_PID"; exit 1' INT TERM

  while true; do
    if [ "$SERVER_PID" -eq -1 ] || ! kill -0 "$SERVER_PID" 2>/dev/null; then
      [ "$SERVER_PID" -ne -1 ] && wait "$SERVER_PID" 2>/dev/null || true
      SERVER_PID=-1
      # echo "running dlv --listen=:2345 --headless=true --api-version=2 exec /usr/local/bin/doltgres -- $@"
      # dlv --listen=:2345 --headless=true --api-version=2 exec /usr/local/bin/doltgres -- "$@" 2>&1 &
      echo "running doltgres $@"
      doltgres "$@" 2>&1 &
      SERVER_PID=$!

    fi

    if is_port_open "0.0.0.0" 5432; then
      note "Doltgres server started."
      return 0
    fi

    local now elapsed
    now=$(date +%s)
    elapsed=$((now - start_time))
    if [ "$elapsed" -ge "$timeout" ]; then
      kill "$SERVER_PID" 2>/dev/null || true
      wait "$SERVER_PID" 2>/dev/null || true
      SERVER_PID=-1
      log_error "Doltgres server failed to start within $timeout seconds"
    fi

    sleep 1
  done
}

write_default_config_yaml() {
    data_dir=$(get_pgdata_var)
    
    if [ -z $data_dir ]; then
        data_dir="/var/lib/doltgres/"
    fi

    if [ ! -d $data_dir ]; then
        mkdir -p $data_dir
    fi
    
    cat <<EOF > config.yaml
log_level: info

encode_logged_query: false

behavior:
  read_only: false
  dolt_transaction_commit: false

listener:
  host: 0.0.0.0
  port: 5432
  read_timeout_millis: 28800000
  write_timeout_millis: 28800000
  allow_cleartext_passwords: false

data_dir: $data_dir

cfg_dir: .doltcfg

privilege_file: .doltcfg/privileges.db

branch_control_file: .doltcfg/branch_control.db
EOF
}
    
# _main is the main entrypoint for the Doltgres Docker container initialization.
_main() {
  check_for_doltgres_binary

  local doltgres_version
  doltgres_version=$(doltgres --version | cut -f3 -d " ")
  note "Entrypoint script for Doltgres Server $doltgres_version starting..."

  declare -g CONFIG_PROVIDED

  CONFIG_PROVIDED=

  # if there is a single yaml provided in /etc/doltgres/servercfg.d directory,
  # it will be used to start the server with --config flag.
  get_config_file_path_if_exists "$SERVER_CONFIG_DIR" "yaml"
  if [ -n "$CONFIG_PROVIDED" ]; then
      set -- "$@" --config="$CONFIG_PROVIDED"
  elif [ -e config.yaml ]; then
      set -- "$@" --config=config.yaml
  else
      echo "generating config.yaml for first run"
      write_default_config_yaml
      cat config.yaml
      set -- "$@" --config=config.yaml
  fi

  note "Starting Doltgres server"
  
  start_server "$@"

  if [[ ! -f $INIT_COMPLETED ]]; then
    if ls /docker-entrypoint-initdb.d/* >/dev/null 2>&1; then
      docker_process_init_files /docker-entrypoint-initdb.d/*
    else
      warn "No files found in /docker-entrypoint-initdb.d/ to process"
    fi
    touch "$INIT_COMPLETED"
  fi

  note "Doltgres running. Ready for connections."
  wait "$SERVER_PID"
}

_main "$@"