#!/bin/bash

DB_BACKUP_MODE="full" # Default mode is 'full'. Options: 'full', 'quick', 'none'

PG_SECRET_NAME=""
PG_SECRET_USER_KEY="username"
PG_SECRET_PASSWORD_KEY="password"

PG_SERVICE_NAME="postgres-pgpool"
PG_POOL_POD_LABEL="app.kubernetes.io/component=pgpool"

BACKUP_POD_IMAGE="postgres:16.4"
BACKUP_POD_NAME="pgdump-debug"

NAMESPACE=""
RELEASE_NAME=""
BACKUP_DIR=""

# Function for initial setup (creating backup directory, etc.)
setup_backup() {
	echo "Setting up backup environment"
	export BACKUP_DIR="./helm-backup-$(date +%Y%m%d-%H%M%S)-${NAMESPACE}"

	mkdir -p "$BACKUP_DIR"
	if [ $? -ne 0 ]; then
		echo "FATAL: Failed to create backup directory $BACKUP_DIR"
		exit 1
	fi

	echo "Backup directory created: $BACKUP_DIR"
}

# Function to back up Kubernetes Secrets and ConfigMaps
backup_k8s_resources() {
	echo "Backing up Kubernetes Resources in namespace ${NAMESPACE}"

	echo "Backing up Secrets..."
	if ! kubectl get secrets -n "$NAMESPACE" -o yaml >"${BACKUP_DIR}/all-secrets.yaml"; then
		echo "ERROR: Failed to back up Secrets."
	else
		echo "Secrets backed up successfully."
	fi

	echo "Backing up ConfigMaps..."
	if ! kubectl get configmaps -n "$NAMESPACE" -o yaml >"${BACKUP_DIR}/all-configmaps.yaml"; then
		echo "ERROR: Failed to back up ConfigMaps."
	else
		echo "ConfigMaps backed up successfully."
	fi
}

# Function to get PostgreSQL credentials from a secret
get_pg_credentials() {
	echo "Retrieving PostgreSQL Credentials from secret ${PG_SECRET_NAME} in namespace ${NAMESPACE}"

	PG_USER=$(kubectl get secret -n "$NAMESPACE" "$PG_SECRET_NAME" -o jsonpath="{.data.$PG_SECRET_USER_KEY}" | base64 --decode)
	PG_PASSWORD=$(kubectl get secret -n "$NAMESPACE" "$PG_SECRET_NAME" -o jsonpath="{.data.$PG_SECRET_PASSWORD_KEY}" | base64 --decode)

	if [ -z "$PG_USER" ] || [ -z "$PG_PASSWORD" ]; then
		echo "FATAL: Failed to decode PostgreSQL user or password from secret $PG_SECRET_NAME in namespace $NAMESPACE."
		return 1
	fi

	export PG_USER
	export PG_PASSWORD
	return 0
}

# Function to launch and wait for the pgdump utility pod
launch_pgdump_pod() {
	echo "Launching pg_dump utility pod $BACKUP_POD_NAME"

	# Ensure the temporary pod is cleaned up if it exists from a previous failed run
	kubectl delete pod "$BACKUP_POD_NAME" -n "$NAMESPACE" --ignore-not-found=true

	kubectl run "$BACKUP_POD_NAME" \
		--namespace="$NAMESPACE" \
		--image="${BACKUP_POD_IMAGE}" \
		--env="PGPASSWORD=${PG_PASSWORD}" \
		--env="PGUSER=${PG_USER}" \
		--restart=Never \
		--wait=false \
		--command -- sleep infinity

	echo "Waiting for $BACKUP_POD_NAME to be Ready (timeout 300s)..."
	if ! kubectl wait --for=condition=Ready pod/"$BACKUP_POD_NAME" --timeout=300s -n "$NAMESPACE"; then
		echo "FATAL: Failed to get $BACKUP_POD_NAME ready within 300 seconds."
		return 1
	fi
	return 0
}

# Function to perform the actual PostgreSQL dump operations
perform_pg_dump() {
	local db_name="$1"
	local output_file="$2"
	local dump_options="$3"

	echo "Backing up database '$db_name' to ${BACKUP_DIR}/$output_file..."

	# Use kubectl exec to run pg_dump inside the utility pod
	if ! kubectl exec -i "$BACKUP_POD_NAME" -n "$NAMESPACE" -- pg_dump $dump_options -U "${PG_USER}" -h "${PG_SERVICE_NAME}" -p 10518 -d "${db_name}" >"${BACKUP_DIR}/$output_file"; then
		echo "ERROR: Failed to back up PostgreSQL database '$db_name'."
		return 1
	fi
	return 0
}

# Function to clean up the temporary pgdump pod
cleanup() {
	echo "Cleaning up $BACKUP_POD_NAME pod in namespace ${NAMESPACE}"
	kubectl delete pod "$BACKUP_POD_NAME" -n "$NAMESPACE" --ignore-not-found=true --wait=false
	echo "Cleanup completed."
}

# Function to perform all PostgreSQL backups
backup_postgresql() {
	echo "Starting PostgreSQL backup..."

	# Find the running PgPool pod
	PGPOOL_POD=$(kubectl get pods -n "$NAMESPACE" -l "$PG_POOL_POD_LABEL" -o jsonpath="{.items[0].metadata.name}")
	if [ -z "$PGPOOL_POD" ]; then
		echo "ERROR: No PostgreSQL pod found with label $PG_POOL_POD_LABEL in namespace $NAMESPACE. Continuing without PG backup."
		return 1
	fi
	echo "Found PostgreSQL pod: $PGPOOL_POD"

	if ! get_pg_credentials; then
		echo "FATAL: Cannot proceed with PostgreSQL backup without credentials."
		return 1
	fi

	if ! launch_pgdump_pod; then
		echo "FATAL: Cannot proceed with PostgreSQL backup without a running utility pod."
		return 1
	fi

	# 1. Globals-only backup
	echo "Backing up globals..."
	if ! kubectl exec -i "$BACKUP_POD_NAME" -n "$NAMESPACE" -- pg_dumpall --globals-only -h "${PG_SERVICE_NAME}" -p 10518 >"${BACKUP_DIR}/db-01-globals.sql"; then
		echo "ERROR: Failed to back up PostgreSQL globals."
	fi

	# 2. Individual database backups
	perform_pg_dump 'atscale' 'db-02-atscale.bin' '-Fc' || true
	perform_pg_dump 'keycloak' 'db-03-keycloak.bin' '-Fc' || true
	perform_pg_dump 'pgwire' 'db-04-pgwire.bin' '-Fc' || true

	# 3. Full dump
	if [ "$DB_BACKUP_MODE" = "full" ]; then
		echo "Backing up PostgreSQL Full Dump..."
		if ! kubectl exec -i "$BACKUP_POD_NAME" -n "$NAMESPACE" -- pg_dumpall -h "${PG_SERVICE_NAME}" -p 10518 >"${BACKUP_DIR}/db-05-full.sql"; then
			echo "ERROR: Failed to back up PostgreSQL Full Dump."
		fi
	else
		echo "Skipping full pg_dumpall as per backup mode '$DB_BACKUP_MODE'."
	fi
}

backup_action() {
	echo "Starting backup for Namespace: ${NAMESPACE} using PG Secret: ${PG_SECRET_NAME}"

	# 1. Setup
	setup_backup

	# 2. Kubernetes Resources Backup
	backup_k8s_resources

	# 3. PostgreSQL Backup
	if [ "$DB_BACKUP_MODE" = "none" ]; then
		echo "Skipping PostgreSQL backup as per --only-secrets option."
		return 0
	fi

	trap cleanup EXIT

	if ! backup_postgresql; then
		echo "WARNING: PostgreSQL backup failed or was skipped."
		exit_code=1
	fi

	# 4. Final message
	echo ""
	echo "****************************************************************"
	echo "Backup Process Finished."
	echo "Backup contents saved to: $BACKUP_DIR"
	echo "****************************************************************"
	echo ""

	cleanup
	return $exit_code
}

get_secrets_action() {
	echo "Encryption: $(kubectl get secret -n "$NAMESPACE" tableau-secret -o json | jq -r .data.ATSCALE_CRYPTO_KEY | base64 -d)"
	echo "--------------------------------------"
	echo "Keycloak Username: $(kubectl get secret -n "$NAMESPACE" keycloak-users-secret -o json | jq -r .data.keycloakAdmin | base64 -d)"
	echo "Keycloak Password: $(kubectl get secret -n "$NAMESPACE" keycloak-users-secret -o json | jq -r .data.keycloakAdminPassword | base64 -d)"
	echo "--------------------------------------"
	echo "Keycloak AtScale Username: $(kubectl get secret -n "$NAMESPACE" keycloak-users-secret -o json | jq -r .data.atscaleAdmin | base64 -d)"
	echo "Keycloak AtScale Password: $(kubectl get secret -n "$NAMESPACE" keycloak-users-secret -o json | jq -r .data.atscaleAdminPassword | base64 -d)"
	echo "--------------------------------------"
	echo "Keycloak API Client Secret: $(kubectl get secret -n "$NAMESPACE" keycloak-clients-secret -o json | jq -r .data.api | base64 -d)"
	echo "Keycloak Engine Client Secret: $(kubectl get secret -n "$NAMESPACE" keycloak-clients-secret -o json | jq -r .data.engine | base64 -d)"
	echo "Keycloak Entitlement Client Secret: $(kubectl get secret -n "$NAMESPACE" keycloak-clients-secret -o json | jq -r .data.entitlement | base64 -d)"
	echo "Keycloak Modeler Client Secret: $(kubectl get secret -n "$NAMESPACE" keycloak-clients-secret -o json | jq -r .data.modeler | base64 -d)"
	echo "Keycloak Public API Client Secret: $(kubectl get secret -n "$NAMESPACE" keycloak-clients-secret -o json | jq -r .data.publicApi | base64 -d)"
	echo "--------------------------------------"
	echo "Postgres User: $(kubectl get secret -n "$NAMESPACE" "postgres-default-user" -o json | jq -r .data.username | base64 -d)"
	echo "Postgres Password: $(kubectl get secret -n "$NAMESPACE" "postgres-default-user" -o json | jq -r .data.password | base64 -d)"
	echo "--------------------------------------"
	echo "PgPool User: $(kubectl get secret -n "$NAMESPACE" "postgres-pgpool-user" -o json | jq -r '.data["admin-username"]' | base64 -d)"
	echo "PgPool Password: $(kubectl get secret -n "$NAMESPACE" "postgres-pgpool-user" -o json | jq -r '.data["admin-password"]' | base64 -d)"
	echo "--------------------------------------"
	echo "Repmgr User: $(kubectl get secret -n "$NAMESPACE" "postgres-default-user" -o json | jq -r '.data["repmgr-username"]' | base64 -d)"
	echo "Repmgr Password: $(kubectl get secret -n "$NAMESPACE" "postgres-default-user" -o json | jq -r '.data["repmgr-password"]' | base64 -d)"
	echo "Repmgr Database: $(kubectl get secret -n "$NAMESPACE" "postgres-default-user" -o json | jq -r '.data["repmgr-database"]' | base64 -d)"
	echo "--------------------------------------"
	echo "Custom Users and Passwords:"
	CUSTOM_USERS=$(kubectl get secret -n "$NAMESPACE" "postgres-custom-users" -o json | jq -r .data.usernames | base64 -d)
	CUSTOM_PASSWORDS=$(kubectl get secret -n "$NAMESPACE" "postgres-custom-users" -o json | jq -r .data.passwords | base64 -d)
	IFS=';' read -r -a user_array <<<"$CUSTOM_USERS"
	IFS=';' read -r -a pass_array <<<"$CUSTOM_PASSWORDS"
	for i in "${!user_array[@]}"; do
		echo "  ${user_array[i]} = ${pass_array[i]}"
	done
	echo "--------------------------------------"
}

restore_action() {
	echo "Starting PostgreSQL restore..."

	# Find the running PgPool pod
	PGPOOL_POD=$(kubectl get pods -n "$NAMESPACE" -l "$PG_POOL_POD_LABEL" -o jsonpath="{.items[0].metadata.name}")
	if [ -z "$PGPOOL_POD" ]; then
		echo "ERROR: No PostgreSQL pod found with label $PG_POOL_POD_LABEL in namespace $NAMESPACE. Continuing without PG backup."
		return 1
	fi
	echo "Found PostgreSQL pod: $PGPOOL_POD"

	if ! get_pg_credentials; then
		echo "FATAL: Cannot proceed with PostgreSQL backup without credentials."
		return 1
	fi

	if ! launch_pgdump_pod; then
		echo "FATAL: Cannot proceed with PostgreSQL backup without a running utility pod."
		return 1
	fi

	echo "Restoring PostgreSQL databases..."

	echo "--Creating databases..."
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d postgres -p 10518 -c 'CREATE DATABASE atscale;'"
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d postgres -p 10518 -c 'CREATE DATABASE keycloak;'"
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d postgres -p 10518 -c 'CREATE DATABASE pgwire;'"

	echo "--Restoring PostgreSQL globals..."
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d postgres -p 10518 " <"$BACKUP_DIR/db-01-globals.sql"

	echo "--Restoring PostgreSQL database 'atscale'..."
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' pg_restore -h ${PGPOOL_SERVICE_NAME} -p 10518  -U ${PG_USER} -d 'atscale' --clean --if-exists" <"$BACKUP_DIR/db-02-atscale.bin"

	echo "--Restoring PostgreSQL database 'keycloak'..."
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' pg_restore -h ${PGPOOL_SERVICE_NAME} -p 10518  -U ${PG_USER} -d 'keycloak' --clean --if-exists" <"$BACKUP_DIR/db-03-keycloak.bin"

	echo "--Restoring PostgreSQL database 'pgwire'..."
	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' pg_restore -h ${PGPOOL_SERVICE_NAME} -p 10518  -U ${PG_USER} -d 'pgwire' --clean --if-exists" <"$BACKUP_DIR/db-04-pgwire.bin"

	echo "--Granting permissions on schemas to the correct users..."

	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d atscale -p 10518 " <<-'EOF'
			\c atscale;
			GRANT ALL PRIVILEGES ON DATABASE atscale TO atscale;
			ALTER DATABASE atscale OWNER TO atscale;
			SET ROLE atscale;
		EOF

	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d keycloak -p 10518 " <<-'EOF'
			\c keycloak;
			GRANT ALL PRIVILEGES ON DATABASE keycloak TO keycloak;
			ALTER DATABASE keycloak OWNER TO keycloak;
			SET ROLE keycloak;
		EOF

	kubectl exec -i -n "$NAMESPACE" pgdump-debug -- \
		bash -c "PGPASSWORD='${PG_PASSWORD}' psql -U ${PG_USER} -h ${PGPOOL_SERVICE_NAME} -d pgwire -p 10518 " <<-'EOF'
			\c pgwire;
			GRANT ALL PRIVILEGES ON DATABASE pgwire TO atscale_metadata;
			ALTER DATABASE pgwire OWNER TO atscale_metadata;
			SET ROLE atscale_metadata;
		EOF

	cleanup
}

usage() {
	echo "Usage: $0 --namespace <k8s_namespace> --pg-secret-name <postgresql_secret_name>"
	echo ""
	echo "Actions:"
	echo "  backup               : Perform a backup of the PostgreSQL database."
	echo "    --db-full          : Perform a full backup including a pg_dumpall, which can take time for large databases."
	echo "    --db-quick         : Perform only globals backup and database dumps."
	echo "    --only-secrets     : Backup only Kubernetes Secrets and ConfigMaps, skipping PostgreSQL backup."
	echo ""
	echo "  get-secrets          : Retrieve and display relevant Kubernetes secrets."
	echo ""
	echo "  restore              : Restore the PostgreSQL database from a backup."
	echo "    --backup-dir <path>          : Path to the directory containing the backup files."
	echo "    --pg-secret-name <name>      : The name of the Kubernetes secret holding PostgreSQL credentials."
	echo "    --pgpool-service-name <name> : The name of the PostgreSQL service (default: postgres-pgpool)."
	echo ""
	echo "Options:"
	echo "  --namespace, -n      : The Kubernetes namespace to backup resources from."
	echo "  --pg-secret-name     : The name of the Kubernetes secret holding PostgreSQL credentials."
	echo "  --help, -h           : Display this help message."
	echo "  --debug, -d          : Enable debug mode (set -x)."
	exit 1
}

parse_command_line_args() {
	while [[ "$#" -gt 0 ]]; do
		case "$1" in
		--namespace | -n)
			NAMESPACE="$2"
			shift 2
			;;
		--pg-secret-name | -s)
			PG_SECRET_NAME="$2"
			shift 2
			;;
		--pgpool-service-name | -p)
			PGPOOL_SERVICE_NAME="$2"
			shift 2
			;;
		--db-full)
			DB_BACKUP_MODE="full"
			shift
			;;
		--db-quick)
			DB_BACKUP_MODE="quick"
			shift
			;;
		--only-secrets)
			DB_BACKUP_MODE="none"
			shift
			;;
		--backup-dir)
			BACKUP_DIR="$2"
			shift 2
			;;
		--help | -h)
			usage
			;;
		--debug | -d)
			set -x
			shift
			;;
		*)
			echo "Error: Unknown option '$1'." >&2
			usage
			;;
		esac
	done
}

main() {
	local exit_code=0

	# Parse Actions
	if [ "$#" -lt 1 ]; then
		echo "Error: No action specified." >&2
		usage
	fi
	ACTION="$1"
	shift

	case "$ACTION" in
	backup)
		parse_command_line_args "$@"

		if [ -z "$NAMESPACE" ]; then
			echo "Error: --namespace is required for backup action." >&2
			exit 1
		fi
		if [ "$DB_BACKUP_MODE" != "none" ] && [ -z "$PG_SECRET_NAME" ]; then
			echo "Error: --pg-secret-name is required for backup action." >&2
			exit 1
		fi

		backup_action || exit_code=1
		;;
	restore)
		parse_command_line_args "$@"

		if [ -z "$NAMESPACE" ] || [ -z "$PG_SECRET_NAME" ] || [ -z "$PGPOOL_SERVICE_NAME" ] || [ -z "$BACKUP_DIR" ]; then
			echo "Error: --namespace is required for restore action." >&2
			echo "Error: --pg-secret-name is required for restore action." >&2
			echo "Error: --pgpool-service-name is required for restore action." >&2
			echo "Error: --backup-dir is required for restore action." >&2
			exit 1
		fi

		restore_action || exit_code=1
		;;
	get-secrets)
		parse_command_line_args "$@"

		if [ -z "$NAMESPACE" ] || [ -z "$PG_SECRET_NAME" ]; then
			echo "Error: --namespace is required for get-secrets action." >&2
			exit 1
		fi

		get_secrets_action || exit_code=1
		;;
	*)
		echo "Error: Unknown action '$ACTION'." >&2
		usage
		;;
	esac

	exit $exit_code
}

main "$@"
