#!/usr/bin/env bash
# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

# Database-only bootstrap for the UmbraTrans control plane.
# This intentionally does not install packages, create application directories,
# install application services, or modify Qwen/Docker/firewall settings.
set -Eeuo pipefail
umask 077

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
CONFIG_FILE="${UMBRA_BOOTSTRAP_CONFIG:-${SCRIPT_DIR}/enterprise-bootstrap.env}"
MODE="check"
DRY_RUN=0
LOCK_DIR="/run/lock/umbracontrol-database.lock"
LOG_FILE="/var/log/umbracontrol-database-bootstrap.log"
PG_SOCKET_DIR="/var/run/postgresql"

UMBRA_DB_NAME="umbracontrol"
UMBRA_DB_HOST="127.0.0.1"
UMBRA_DB_PORT="5432"
UMBRA_BOOTSTRAP_ADMIN_EMAIL="umbratools@gmail.com"
UMBRA_ORGANIZATION_NAME="Umbra"
UMBRA_ORGANIZATION_SLUG="umbra"
UMBRA_PG_SUPERUSER="postgres"
UMBRA_CREDENTIALS_FILE="/etc/umbracontrol/db-credentials.env"
UMBRA_BOOTSTRAP_TOKEN_FILE="/root/umbra-enterprise-bootstrap.token"
UMBRA_SERVICE_GROUP="umbracontrol"
UMBRA_OWNER_ROLE="umbracontrol_owner"
UMBRA_MIGRATOR_ROLE="umbracontrol_migrator"
UMBRA_API_ROLE="umbracontrol_api"
UMBRA_WORKER_ROLE="umbracontrol_worker"
UMBRA_READONLY_ROLE="umbracontrol_readonly"
UMBRA_BACKUP_ROLE="umbracontrol_backup"

usage() {
    cat <<'EOF'
Usage: bootstrap_umbracontrol_database.sh [options]

Options:
  --check       Read-only PostgreSQL and prerequisite inspection
  --dry-run     Show database actions without changing the host
  --apply       Create/finish the database, roles, schema, and bootstrap admin
  --verify      Verify the database, schema, RLS, and bootstrap records
  --config FILE Read the existing non-secret enterprise-bootstrap.env
  -h, --help    Show this help

This script is intentionally database-only. It does not run apt, install uv,
create application directories, start Qwen, or change firewall rules.
EOF
}

while (($#)); do
    case "$1" in
        --check) MODE="check"; shift ;;
        --dry-run) MODE="apply"; DRY_RUN=1; shift ;;
        --apply) MODE="apply"; DRY_RUN=0; shift ;;
        --verify) MODE="verify"; shift ;;
        --config) CONFIG_FILE="${2:?--config requires a file}"; shift 2 ;;
        -h|--help) usage; exit 0 ;;
        *) printf 'Unknown option: %s\n' "$1" >&2; usage >&2; exit 2 ;;
    esac
done

log() {
    local line="[$(date -Is)] $*"
    printf '%s\n' "$line"
    if [[ -f "$LOG_FILE" && "$(id -u)" -eq 0 ]]; then
        printf '%s\n' "$line" >> "$LOG_FILE"
    fi
}

die() {
    printf 'ERROR: %s\n' "$*" >&2
    exit 1
}

require_root() {
    [[ "$(id -u)" -eq 0 ]] || die "This operation requires root. Use --check for read-only inspection."
}

require_command() {
    command -v "$1" >/dev/null 2>&1 || die "Required command is missing: $1"
}

load_config() {
    [[ -f "$CONFIG_FILE" ]] || {
        [[ "$MODE" == "check" || "$MODE" == "verify" ]] && return 0
        die "Configuration file not found: $CONFIG_FILE"
    }

    local line key value
    while IFS= read -r line || [[ -n "$line" ]]; do
        line="${line#"${line%%[![:space:]]*}"}"
        [[ -z "$line" || "${line:0:1}" == "#" ]] && continue
        [[ "$line" == *=* ]] || die "Invalid configuration line: $line"
        key="${line%%=*}"
        value="${line#*=}"
        key="${key//[[:space:]]/}"
        value="${value%\"}"; value="${value#\"}"
        value="${value%\'}"; value="${value#\'}"
        case "$key" in
            UMBRA_DB_NAME|UMBRA_DB_HOST|UMBRA_DB_PORT|UMBRA_BOOTSTRAP_ADMIN_EMAIL|UMBRA_ORGANIZATION_NAME|UMBRA_ORGANIZATION_SLUG)
                printf -v "$key" '%s' "$value"
                ;;
            UMBRA_*)
                : # Ignore unrelated non-secret enterprise settings.
                ;;
            *)
                die "Unsupported configuration key: $key"
                ;;
        esac
    done < "$CONFIG_FILE"
}

validate_config() {
    [[ "$UMBRA_DB_NAME" =~ ^[a-z_][a-z0-9_]*$ ]] || die "Invalid UMBRA_DB_NAME"
    [[ "$UMBRA_DB_HOST" =~ ^[A-Za-z0-9._:-]+$ ]] || die "Invalid UMBRA_DB_HOST"
    [[ "$UMBRA_DB_PORT" =~ ^[0-9]+$ && "$UMBRA_DB_PORT" -ge 1 && "$UMBRA_DB_PORT" -le 65535 ]] || die "Invalid UMBRA_DB_PORT"
    [[ "$UMBRA_BOOTSTRAP_ADMIN_EMAIL" == *@*.* ]] || die "Invalid administrator email"
    [[ "$UMBRA_ORGANIZATION_SLUG" =~ ^[a-z0-9][a-z0-9-]*$ ]] || die "Invalid organization slug"
    [[ -d "$PG_SOCKET_DIR" ]] || die "PostgreSQL socket directory is missing: $PG_SOCKET_DIR"
}

run() {
    if ((DRY_RUN)); then
        printf '[dry-run]'
        printf ' %q' "$@"
        printf '\n'
    else
        "$@"
    fi
}

random_secret() {
    openssl rand -hex 32
}

read_secret() {
    local key="$1"
    [[ -f "$UMBRA_CREDENTIALS_FILE" ]] || return 1
    awk -F= -v key="$key" '$1 == key {print substr($0, index($0, "=") + 1); exit}' "$UMBRA_CREDENTIALS_FILE"
}

write_secret() {
    local key="$1" value="$2" temporary
    if ((DRY_RUN)); then
        printf '[dry-run] store generated database secret %s\n' "$key"
        return
    fi
    install -d -m 0750 -o root -g "$UMBRA_SERVICE_GROUP" "$(dirname "$UMBRA_CREDENTIALS_FILE")"
    touch "$UMBRA_CREDENTIALS_FILE"
    chown root:"$UMBRA_SERVICE_GROUP" "$UMBRA_CREDENTIALS_FILE"
    chmod 0640 "$UMBRA_CREDENTIALS_FILE"
    temporary="$(mktemp "${UMBRA_CREDENTIALS_FILE}.XXXXXX")"
    awk -F= -v key="$key" -v value="$value" '
        $1 == key { print key "=" value; found = 1; next }
        { print }
        END { if (!found) print key "=" value }
    ' "$UMBRA_CREDENTIALS_FILE" > "$temporary"
    chown root:"$UMBRA_SERVICE_GROUP" "$temporary"
    chmod 0640 "$temporary"
    mv -f "$temporary" "$UMBRA_CREDENTIALS_FILE"
}

psql_as_postgres() {
    runuser -u "$UMBRA_PG_SUPERUSER" -- env \
        -u PGHOST -u PGPORT -u PGDATABASE -u PGUSER -u PGPASSWORD \
        -u PGPASSWORDFILE -u PGSERVICE -u PGSERVICEFILE \
        PGOPTIONS='-c log_statement=none' \
        psql -h "$PG_SOCKET_DIR" -p "$UMBRA_DB_PORT" -U "$UMBRA_PG_SUPERUSER" \
        -v ON_ERROR_STOP=1 "$@"
}

role_exists() {
    local role="$1"
    psql_as_postgres postgres -tAc "SELECT 1 FROM pg_roles WHERE rolname = '$role'" | grep -q 1
}

ensure_role() {
    local role="$1" login="$2" password_key="$3" password=""
    if ((DRY_RUN)); then
        printf '[dry-run] ensure PostgreSQL role %s (%s)\n' "$role" "$login"
        return
    fi
    if role_exists "$role"; then
        if [[ "$login" == "LOGIN" ]]; then
            printf 'ALTER ROLE "%s" LOGIN;\n' "$role" | psql_as_postgres postgres >/dev/null
        else
            printf 'ALTER ROLE "%s" NOLOGIN;\n' "$role" | psql_as_postgres postgres >/dev/null
        fi
        log "PostgreSQL role already exists: $role"
        return
    fi

    if [[ "$login" == "LOGIN" ]]; then
        password="$(read_secret "$password_key" 2>/dev/null || true)"
        [[ -n "$password" ]] || password="$(random_secret)"
        write_secret "$password_key" "$password"
        printf 'CREATE ROLE "%s" LOGIN PASSWORD '\''%s'\'';\n' "$role" "$password" \
            | psql_as_postgres postgres >/dev/null
    else
        printf 'CREATE ROLE "%s" NOLOGIN;\n' "$role" | psql_as_postgres postgres >/dev/null
    fi
    log "PostgreSQL role created: $role"
}

ensure_roles() {
    log "Ensuring PostgreSQL roles"
    ensure_role "$UMBRA_OWNER_ROLE" NOLOGIN OWNER_UNUSED
    ensure_role "$UMBRA_MIGRATOR_ROLE" LOGIN UMBRACONTROL_MIGRATOR_PASSWORD
    ensure_role "$UMBRA_API_ROLE" LOGIN UMBRACONTROL_API_PASSWORD
    ensure_role "$UMBRA_WORKER_ROLE" LOGIN UMBRACONTROL_WORKER_PASSWORD
    ensure_role "$UMBRA_READONLY_ROLE" LOGIN UMBRACONTROL_READONLY_PASSWORD
    ensure_role "$UMBRA_BACKUP_ROLE" LOGIN UMBRACONTROL_BACKUP_PASSWORD
}

database_exists() {
    psql_as_postgres postgres -tAc "SELECT 1 FROM pg_database WHERE datname = '$UMBRA_DB_NAME'" | grep -q 1
}

ensure_database() {
    log "Ensuring database: $UMBRA_DB_NAME"
    if ((DRY_RUN)); then
        printf '[dry-run] create database if absent, owner=%s\n' "$UMBRA_OWNER_ROLE"
        return
    fi
    if ! database_exists; then
        runuser -u "$UMBRA_PG_SUPERUSER" -- createdb --owner="$UMBRA_OWNER_ROLE" "$UMBRA_DB_NAME"
        log "Database created"
    else
        log "Database already exists"
    fi
    psql_as_postgres postgres -c "REVOKE ALL ON DATABASE \"$UMBRA_DB_NAME\" FROM PUBLIC" >/dev/null
    psql_as_postgres postgres -c "GRANT CONNECT ON DATABASE \"$UMBRA_DB_NAME\" TO \"$UMBRA_MIGRATOR_ROLE\", \"$UMBRA_API_ROLE\", \"$UMBRA_WORKER_ROLE\", \"$UMBRA_READONLY_ROLE\", \"$UMBRA_BACKUP_ROLE\"" >/dev/null
}

apply_schema() {
    local schema_file="${SCRIPT_DIR}/postgres/001_control_plane.sql"
    [[ -f "$schema_file" ]] || die "Schema file is missing: $schema_file"
    log "Applying identity, tenancy, role, audit, and RLS schema"
    if ((DRY_RUN)); then
        printf '[dry-run] stream %s into PostgreSQL database %s\n' "$schema_file" "$UMBRA_DB_NAME"
        return
    fi
    /bin/cat -- "$schema_file" | psql_as_postgres "$UMBRA_DB_NAME" >/dev/null
    log "Database schema applied"
}

create_bootstrap_admin() {
    local token token_hash
    if [[ -f "$UMBRA_BOOTSTRAP_TOKEN_FILE" ]]; then
        token="$(cat "$UMBRA_BOOTSTRAP_TOKEN_FILE")"
    else
        token="$(random_secret)"
        if ((DRY_RUN)); then
            printf '[dry-run] create root-only bootstrap token\n'
            return
        fi
        printf '%s\n' "$token" > "$UMBRA_BOOTSTRAP_TOKEN_FILE"
        chown root:root "$UMBRA_BOOTSTRAP_TOKEN_FILE"
        chmod 0600 "$UMBRA_BOOTSTRAP_TOKEN_FILE"
    fi
    [[ "$token" =~ ^[a-f0-9]{64}$ ]] || die "Bootstrap token has an unexpected format"
    token_hash="$(printf '%s' "$token" | sha256sum | awk '{print $1}')"

    log "Ensuring bootstrap administrator"
    if ((DRY_RUN)); then
        printf '[dry-run] create administrator %s and organization %s\n' "$UMBRA_BOOTSTRAP_ADMIN_EMAIL" "$UMBRA_ORGANIZATION_NAME"
        return
    fi
    psql_as_postgres "$UMBRA_DB_NAME" \
        -v admin_email="$UMBRA_BOOTSTRAP_ADMIN_EMAIL" \
        -v token_hash="$token_hash" \
        -v organization_name="$UMBRA_ORGANIZATION_NAME" \
        -v organization_slug="$UMBRA_ORGANIZATION_SLUG" <<'SQL' >/dev/null
INSERT INTO users (email, display_name, auth_provider, auth_subject, status, mfa_required)
VALUES (:'admin_email', :'admin_email', 'bootstrap', 'bootstrap:' || gen_random_uuid()::text, 'invited', true)
ON CONFLICT (email) DO UPDATE
    SET display_name = EXCLUDED.display_name,
        auth_provider = 'bootstrap',
        mfa_required = true,
        updated_at = now();

INSERT INTO organizations (name, slug)
VALUES (:'organization_name', :'organization_slug')
ON CONFLICT (slug) DO UPDATE SET name = EXCLUDED.name;

INSERT INTO user_roles (user_id, role_id, scope, organization_id, project_id, created_by)
SELECT u.id, r.id, 'platform', NULL, NULL, u.id
FROM users u
JOIN roles r ON r.scope = 'platform' AND r.name = 'platform_admin'
WHERE u.email = :'admin_email'
ON CONFLICT DO NOTHING;

INSERT INTO user_roles (user_id, role_id, scope, organization_id, project_id, created_by)
SELECT u.id, r.id, 'organization', o.id, NULL, u.id
FROM users u
JOIN organizations o ON o.slug = :'organization_slug'
JOIN roles r ON r.scope = 'organization' AND r.name = 'organization_admin'
WHERE u.email = :'admin_email'
ON CONFLICT DO NOTHING;

INSERT INTO bootstrap_tokens (user_id, token_hash, expires_at)
SELECT id, :'token_hash', now() + interval '24 hours'
FROM users
WHERE email = :'admin_email'
ON CONFLICT (token_hash) DO NOTHING;
SQL
    log "Bootstrap administrator is ready"
}

check_mode() {
    log "Database-only read-only inspection"
    printf 'psql: '; command -v psql 2>/dev/null || echo 'missing'
    printf 'PostgreSQL: '; command -v pg_lsclusters >/dev/null 2>&1 && pg_lsclusters || echo 'not installed'
    printf 'service: '; systemctl is-active postgresql 2>/dev/null || true
    printf 'socket: '; test -S "$PG_SOCKET_DIR/.s.PGSQL.$UMBRA_DB_PORT" && echo ready || echo missing
    printf 'database socket: '; pg_isready -h "$PG_SOCKET_DIR" -p "$UMBRA_DB_PORT" -d "$UMBRA_DB_NAME" 2>&1 || true
    printf 'roles: '; getent group "$UMBRA_SERVICE_GROUP" 2>/dev/null || echo 'service group absent'
    log "No changes made"
}

apply_mode() {
    if ((DRY_RUN == 0)); then
        require_root
        mkdir "$LOCK_DIR" 2>/dev/null || die "Another database bootstrap is running"
        trap 'rmdir "$LOCK_DIR" 2>/dev/null || true' EXIT
        touch "$LOG_FILE"
        chmod 0640 "$LOG_FILE"
    else
        log "Database-only bootstrap dry-run"
    fi
    ensure_roles
    ensure_database
    apply_schema
    create_bootstrap_admin
    log "Database-only bootstrap completed"
}

verify_mode() {
    require_root
    log "Verifying database-only bootstrap"
    pg_isready -h "$UMBRA_DB_HOST" -p "$UMBRA_DB_PORT" -d "$UMBRA_DB_NAME"
    psql_as_postgres "$UMBRA_DB_NAME" -v ON_ERROR_STOP=1 \
        -c 'SELECT version, applied_at FROM schema_migrations ORDER BY version' \
        -c 'SELECT count(*) AS organizations FROM organizations' \
        -c 'SELECT count(*) AS users FROM users' \
        -c "SELECT count(*) AS rls_tables FROM pg_class WHERE relrowsecurity AND relname IN ('organizations','projects','users','user_roles','invitations','sessions','mfa_methods','bootstrap_tokens','audit_events')"
    log "Database verification completed"
}

main() {
    load_config
    validate_config
    require_command openssl
    require_command systemctl
    if [[ "$MODE" == "apply" || "$MODE" == "verify" ]]; then
        require_command psql
        require_command pg_isready
        require_command runuser
    fi
    case "$MODE" in
        check) check_mode ;;
        apply) apply_mode ;;
        verify) verify_mode ;;
        *) die "Unknown mode: $MODE" ;;
    esac
}

main