#!/usr/bin/env bash
#
# Copyright (c) 2025, Cloudera, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

#
# One-time setup script: creates CMA databases and users in a PostgreSQL instance.
#
# Usage:
#   ./setup-cma-postgres.sh --host <pg-host> --port <pg-port> --admin-user <superuser> --admin-password <password>
#   ./setup-cma-postgres.sh --host <pg-host> --admin-user postgres --admin-password <password>
#
# The script creates:
#   - Database 'cma' with user 'cma'
#   - Database 'cma-auth' with user 'cma' (shared user by default, or separate 'cma-auth' user with --auth-user)
#   - Database 'cma-agent' with user 'cma-agent'
#
# After running this script, configure CMA in Cloudera Manager:
#   Database Type = postgres
#   Database Host = <pg-host>
#   Database User = cma          (for CMA Master)
#   Database User = cma-agent    (for CMA Agent)
#

set -euo pipefail

PG_HOST=""
PG_PORT="5432"
PG_ADMIN_USER=""
PG_ADMIN_PASS=""
CMA_DB_USER="cma"
CMA_DB_PASS=""
CMA_AUTH_DB_USER=""
CMA_AUTH_DB_PASS=""
CMA_AGENT_DB_USER="cma-agent"
CMA_AGENT_DB_PASS=""

usage() {
  echo "Usage: $(basename "$0") --host <pg-host> --admin-user <superuser> --admin-password <password> [options]"
  echo ""
  echo "Creates CMA databases and users in a PostgreSQL instance."
  echo ""
  echo "Required:"
  echo "  --host              PostgreSQL hostname or IP"
  echo "  --admin-user        PostgreSQL superuser (e.g. postgres)"
  echo "  --admin-password    PostgreSQL superuser password"
  echo ""
  echo "Optional:"
  echo "  --port              PostgreSQL port (default: 5432)"
  echo "  --cma-user          CMA Master database user (default: cma)"
  echo "  --cma-password      CMA Master database user password (auto-generated if omitted)"
  echo "  --auth-user         CMA Auth database user (default: same as --cma-user)"
  echo "  --auth-password     CMA Auth database user password (default: same as --cma-password)"
  echo "  --agent-user        CMA Agent database user (default: cma-agent)"
  echo "  --agent-password    CMA Agent database user password (auto-generated if omitted)"
  echo "  -h, --help          Show this help message"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
  case $1 in
    --host)            PG_HOST="$2"; shift 2 ;;
    --host=*)          PG_HOST="${1#*=}"; shift ;;
    --port)            PG_PORT="$2"; shift 2 ;;
    --port=*)          PG_PORT="${1#*=}"; shift ;;
    --admin-user)      PG_ADMIN_USER="$2"; shift 2 ;;
    --admin-user=*)    PG_ADMIN_USER="${1#*=}"; shift ;;
    --admin-password)  PG_ADMIN_PASS="$2"; shift 2 ;;
    --admin-password=*) PG_ADMIN_PASS="${1#*=}"; shift ;;
    --cma-user)        CMA_DB_USER="$2"; shift 2 ;;
    --cma-user=*)      CMA_DB_USER="${1#*=}"; shift ;;
    --cma-password)    CMA_DB_PASS="$2"; shift 2 ;;
    --cma-password=*)  CMA_DB_PASS="${1#*=}"; shift ;;
    --auth-user)       CMA_AUTH_DB_USER="$2"; shift 2 ;;
    --auth-user=*)     CMA_AUTH_DB_USER="${1#*=}"; shift ;;
    --auth-password)   CMA_AUTH_DB_PASS="$2"; shift 2 ;;
    --auth-password=*) CMA_AUTH_DB_PASS="${1#*=}"; shift ;;
    --agent-user)      CMA_AGENT_DB_USER="$2"; shift 2 ;;
    --agent-user=*)    CMA_AGENT_DB_USER="${1#*=}"; shift ;;
    --agent-password)  CMA_AGENT_DB_PASS="$2"; shift 2 ;;
    --agent-password=*) CMA_AGENT_DB_PASS="${1#*=}"; shift ;;
    -h|--help)         usage; exit 0 ;;
    *)                 echo "Unknown option: $1"; usage; exit 1 ;;
  esac
done

# Validate required args
if [ -z "$PG_HOST" ] || [ -z "$PG_ADMIN_USER" ] || [ -z "$PG_ADMIN_PASS" ]; then
  echo "ERROR: --host, --admin-user, and --admin-password are required."
  echo ""
  usage
  exit 1
fi

# Auto-generate passwords if not provided
generate_password() {
  openssl rand -base64 24 2>/dev/null || head -c 32 /dev/urandom | base64 | tr -d '/+=' | head -c 24
}

if [ -z "$CMA_DB_PASS" ]; then
  CMA_DB_PASS=$(generate_password)
fi
# Auth user defaults to CMA Master user when not specified
if [ -z "$CMA_AUTH_DB_USER" ]; then
  CMA_AUTH_DB_USER="$CMA_DB_USER"
fi
if [ -z "$CMA_AUTH_DB_PASS" ]; then
  CMA_AUTH_DB_PASS="$CMA_DB_PASS"
fi
if [ -z "$CMA_AGENT_DB_PASS" ]; then
  CMA_AGENT_DB_PASS=$(generate_password)
fi

# Verify psql is available
if ! command -v psql &>/dev/null; then
  echo "ERROR: psql command not found. Install the PostgreSQL client:"
  echo "  yum install postgresql  # RHEL/CentOS"
  echo "  apt install postgresql-client  # Debian/Ubuntu"
  exit 1
fi

export PGPASSWORD="$PG_ADMIN_PASS"
trap 'unset PGPASSWORD' EXIT
PSQL=(psql -h "${PG_HOST}" -p "${PG_PORT}" -U "${PG_ADMIN_USER}" -v ON_ERROR_STOP=1)

# Test connection
echo "Connecting to PostgreSQL at ${PG_HOST}:${PG_PORT} as ${PG_ADMIN_USER}..."
if ! "${PSQL[@]}" -d postgres -c "SELECT 1" &>/dev/null; then
  echo "ERROR: Cannot connect to PostgreSQL at ${PG_HOST}:${PG_PORT} as ${PG_ADMIN_USER}."
  echo "Verify the host, port, and credentials."
  exit 1
fi
echo "Connection verified."

# Escape single quotes for safe SQL string interpolation (' -> '')
sql_escape() {
  echo "${1//\'/\'\'}"
}

# Create user if not exists, or update password
create_user() {
  local db_user="$1"
  local db_pass
  db_pass=$(sql_escape "$2")
  if "${PSQL[@]}" -d postgres -tAc "SELECT 1 FROM pg_roles WHERE rolname = '${db_user}'" | grep -q 1; then
    echo "User '${db_user}' already exists, updating password."
    "${PSQL[@]}" -d postgres -c "ALTER USER \"${db_user}\" WITH PASSWORD '${db_pass}';"
  else
    echo "Creating user '${db_user}'..."
    "${PSQL[@]}" -d postgres -c "CREATE USER \"${db_user}\" WITH PASSWORD '${db_pass}';"
  fi
}

# Create database if not exists, grant privileges
create_db() {
  local db_name="$1"
  local db_owner="$2"
  if "${PSQL[@]}" -d postgres -tAc "SELECT 1 FROM pg_database WHERE datname = '${db_name}'" | grep -q 1; then
    echo "Database '${db_name}' already exists, skipping creation."
  else
    echo "Creating database '${db_name}' owned by '${db_owner}'..."
    "${PSQL[@]}" -d postgres -c "CREATE DATABASE \"${db_name}\" OWNER \"${db_owner}\";"
  fi
  "${PSQL[@]}" -d postgres -c "GRANT ALL PRIVILEGES ON DATABASE \"${db_name}\" TO \"${db_owner}\";"
  "${PSQL[@]}" -d "${db_name}" -c "GRANT ALL ON SCHEMA public TO \"${db_owner}\";" 2>/dev/null || true
}

# Create CMA Master user and database
create_user "$CMA_DB_USER" "$CMA_DB_PASS"
create_db "cma" "$CMA_DB_USER"

# Create CMA Auth user and database (may be same user as Master)
if [ "$CMA_AUTH_DB_USER" != "$CMA_DB_USER" ]; then
  create_user "$CMA_AUTH_DB_USER" "$CMA_AUTH_DB_PASS"
fi
create_db "cma-auth" "$CMA_AUTH_DB_USER"

# Create CMA Agent user and database
create_user "$CMA_AGENT_DB_USER" "$CMA_AGENT_DB_PASS"
create_db "cma-agent" "$CMA_AGENT_DB_USER"

mask_password() {
  local pass="$1"
  if [ ${#pass} -le 4 ]; then
    echo "****"
  else
    echo "${pass:0:2}$(printf '*%.0s' $(seq 1 $((${#pass} - 4))))${pass: -2}"
  fi
}

echo ""
echo "======================================"
echo "CMA PostgreSQL setup complete!"
echo "======================================"
echo ""
echo "CMA Master database:"
echo "  Database: cma"
echo "  User:     ${CMA_DB_USER}"
echo "  Password: $(mask_password "${CMA_DB_PASS}")"
echo ""
echo "CMA Auth database:"
echo "  Database: cma-auth"
echo "  User:     ${CMA_AUTH_DB_USER}"
echo "  Password: $(mask_password "${CMA_AUTH_DB_PASS}")"
if [ "$CMA_AUTH_DB_USER" = "$CMA_DB_USER" ]; then
  echo "  (shared with CMA Master user)"
fi
echo ""
echo "CMA Agent database:"
echo "  Database: cma-agent"
echo "  User:     ${CMA_AGENT_DB_USER}"
echo "  Password: $(mask_password "${CMA_AGENT_DB_PASS}")"
echo ""
echo "Configure in Cloudera Manager:"
echo "  CMA Master service:"
echo "    Database Type          = postgres"
echo "    Database Host          = ${PG_HOST}"
echo "    Database Port          = ${PG_PORT}"
echo "    Database User          = ${CMA_DB_USER}"
echo "    Database Password      = <use the password shown above>"
if [ "$CMA_AUTH_DB_USER" != "$CMA_DB_USER" ]; then
  echo "    Auth Database User     = ${CMA_AUTH_DB_USER}"
  echo "    Auth Database Password = <use the password shown above>"
fi
echo ""
echo "  CMA Agent service:"
echo "    Database Type          = postgres"
echo "    Database Host          = ${PG_HOST}"
echo "    Database Port          = ${PG_PORT}"
echo "    Database User          = ${CMA_AGENT_DB_USER}"
echo "    Database Password      = <use the password shown above>"
echo ""
echo "IMPORTANT: Save these passwords securely. They will not be shown again."
