#!/usr/bin/env bash
set -Eeuo pipefail

umask 077

readonly schema_name="rag"
readonly marker_path="${PGDATA}/.rag-bootstrap-complete"

require_file() {
  local path="${1:-}"
  local label="$2"

  if [[ -z "$path" || ! -r "$path" ]]; then
    printf 'RAG bootstrap error: %s file is missing or unreadable\n' "$label" >&2
    return 1
  fi
}

read_secret() {
  local path="$1"
  local label="$2"
  local value

  require_file "$path" "$label"
  value="$(<"$path")"
  if [[ -z "$value" ]]; then
    printf 'RAG bootstrap error: %s must not be empty\n' "$label" >&2
    return 1
  fi
  if [[ "$value" == *$'\n'* || "$value" == *$'\r'* ]]; then
    printf 'RAG bootstrap error: %s must contain exactly one line\n' "$label" >&2
    return 1
  fi
  if (( ${#value} < 16 )); then
    printf 'RAG bootstrap error: %s must contain at least 16 characters\n' "$label" >&2
    return 1
  fi

  printf '%s' "$value"
}

require_identifier() {
  local value="${1:-}"
  local label="$2"

  if [[ ! "$value" =~ ^[A-Za-z_][A-Za-z0-9_]{0,62}$ ]]; then
    printf 'RAG bootstrap error: %s is not a valid PostgreSQL identifier\n' "$label" >&2
    return 1
  fi
}

readonly bootstrap_user="${POSTGRES_USER:?POSTGRES_USER is required}"
readonly database_name="${POSTGRES_DB:?POSTGRES_DB is required}"
readonly migrator_user="${POSTGRES_MIGRATOR_USER:?POSTGRES_MIGRATOR_USER is required}"
readonly app_user="${POSTGRES_APP_USER:?POSTGRES_APP_USER is required}"

require_identifier "$bootstrap_user" "POSTGRES_USER"
require_identifier "$database_name" "POSTGRES_DB"
require_identifier "$migrator_user" "POSTGRES_MIGRATOR_USER"
require_identifier "$app_user" "POSTGRES_APP_USER"

if [[ "$bootstrap_user" == "$migrator_user" ||
      "$bootstrap_user" == "$app_user" ||
      "$migrator_user" == "$app_user" ]]; then
  printf 'RAG bootstrap error: bootstrap, migrator, and app roles must be distinct\n' >&2
  exit 1
fi

# The official postgres entrypoint resolves POSTGRES_PASSWORD_FILE and then
# unsets that path before it invokes /docker-entrypoint-initdb.d scripts.
bootstrap_password="${POSTGRES_PASSWORD:-}"
if [[ -z "$bootstrap_password" || ${#bootstrap_password} -lt 16 ||
      "$bootstrap_password" == *$'\n'* || "$bootstrap_password" == *$'\r'* ]]; then
  printf 'RAG bootstrap error: resolved bootstrap password is missing or too short\n' >&2
  exit 1
fi
migrator_password="$(read_secret "${POSTGRES_MIGRATOR_PASSWORD_FILE:?POSTGRES_MIGRATOR_PASSWORD_FILE is required}" "migrator password")"
app_password="$(read_secret "${POSTGRES_APP_PASSWORD_FILE:?POSTGRES_APP_PASSWORD_FILE is required}" "app password")"

if [[ "$bootstrap_password" == "$migrator_password" ||
      "$bootstrap_password" == "$app_password" ||
      "$migrator_password" == "$app_password" ]]; then
  printf 'RAG bootstrap error: bootstrap, migrator, and app passwords must be distinct\n' >&2
  exit 1
fi

export RAG_MIGRATOR_PASSWORD="$migrator_password"
export RAG_APP_PASSWORD="$app_password"

psql \
  --set=ON_ERROR_STOP=1 \
  --username "$bootstrap_user" \
  --dbname "$database_name" \
  --set=database_name="$database_name" \
  --set=migrator_user="$migrator_user" \
  --set=app_user="$app_user" <<'SQL'
\getenv migrator_password RAG_MIGRATOR_PASSWORD
\getenv app_password RAG_APP_PASSWORD

BEGIN;

CREATE EXTENSION IF NOT EXISTS vector WITH SCHEMA public;

CREATE ROLE :"migrator_user"
  LOGIN
  PASSWORD :'migrator_password'
  NOSUPERUSER
  NOCREATEDB
  NOCREATEROLE
  NOINHERIT
  NOBYPASSRLS;

CREATE ROLE :"app_user"
  LOGIN
  PASSWORD :'app_password'
  NOSUPERUSER
  NOCREATEDB
  NOCREATEROLE
  NOINHERIT
  NOBYPASSRLS;

REVOKE ALL ON DATABASE :"database_name" FROM PUBLIC;
GRANT CONNECT ON DATABASE :"database_name" TO :"migrator_user", :"app_user";

REVOKE CREATE ON SCHEMA public FROM PUBLIC;
GRANT USAGE ON SCHEMA public TO :"migrator_user", :"app_user";

CREATE SCHEMA rag AUTHORIZATION :"migrator_user";
GRANT USAGE, CREATE ON SCHEMA rag TO :"migrator_user";
GRANT USAGE ON SCHEMA rag TO :"app_user";

ALTER ROLE :"migrator_user" IN DATABASE :"database_name"
  SET search_path = rag, public;
ALTER ROLE :"app_user" IN DATABASE :"database_name"
  SET search_path = rag, public;

ALTER DEFAULT PRIVILEGES FOR ROLE :"migrator_user" IN SCHEMA rag
  GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO :"app_user";
ALTER DEFAULT PRIVILEGES FOR ROLE :"migrator_user" IN SCHEMA rag
  GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO :"app_user";
ALTER DEFAULT PRIVILEGES FOR ROLE :"migrator_user" IN SCHEMA rag
  REVOKE EXECUTE ON FUNCTIONS FROM PUBLIC;
ALTER DEFAULT PRIVILEGES FOR ROLE :"migrator_user" IN SCHEMA rag
  GRANT EXECUTE ON FUNCTIONS TO :"app_user";
ALTER DEFAULT PRIVILEGES FOR ROLE :"migrator_user" IN SCHEMA rag
  GRANT USAGE ON TYPES TO :"app_user";

COMMIT;
SQL

unset RAG_MIGRATOR_PASSWORD RAG_APP_PASSWORD

marker_tmp="$(mktemp "${PGDATA}/.rag-bootstrap-complete.tmp.XXXXXX")"
trap 'rm -f "$marker_tmp"' EXIT
printf 'bootstrap_version=1\nschema=%s\n' "$schema_name" >"$marker_tmp"
chmod 0600 "$marker_tmp"
mv -f "$marker_tmp" "$marker_path"
trap - EXIT

unset bootstrap_password migrator_password app_password
printf 'RAG PostgreSQL bootstrap completed successfully.\n'