#!/bin/sh
set -e

# OpenBao Auto-Init Script
# Auto-initializes, unseals, and configures OpenBao on first run
# Idempotent - safe to run multiple times

INIT_DIR="/openbao/init"
UNSEAL_KEY_FILE="${INIT_DIR}/unseal-key"
ROOT_TOKEN_FILE="${INIT_DIR}/root-token"
APPROLE_CREDS_FILE="${INIT_DIR}/approle-credentials"

VAULT_ADDR="http://openbao:8200"
export VAULT_ADDR

# Ensure init directory exists
mkdir -p "${INIT_DIR}"

echo "=== OpenBao Auto-Init Script ==="
echo "Vault Address: ${VAULT_ADDR}"

# Wait for OpenBao to be ready
echo "Waiting for OpenBao API to be available..."
MAX_RETRIES=30
RETRY_COUNT=0
while ! wget -q -O- "${VAULT_ADDR}/v1/sys/init" >/dev/null 2>&1; do
  RETRY_COUNT=$((RETRY_COUNT + 1))
  if [ ${RETRY_COUNT} -ge ${MAX_RETRIES} ]; then
    echo "ERROR: OpenBao API not available after ${MAX_RETRIES} attempts"
    exit 1
  fi
  echo "Waiting for OpenBao... (${RETRY_COUNT}/${MAX_RETRIES})"
  sleep 2
done
echo "OpenBao API is available"

# Check initialization status
INIT_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/init" 2>/dev/null || echo '{"initialized":false}')
IS_INITIALIZED=$(echo "${INIT_STATUS}" | grep -o '"initialized":[^,}]*' | cut -d':' -f2)

if [ "${IS_INITIALIZED}" = "true" ]; then
  echo "OpenBao is already initialized"

  # Check if sealed
  SEAL_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/seal-status" 2>/dev/null)
  IS_SEALED=$(echo "${SEAL_STATUS}" | grep -o '"sealed":[^,}]*' | cut -d':' -f2)

  if [ "${IS_SEALED}" = "true" ]; then
    echo "OpenBao is sealed - unsealing..."

    if [ ! -f "${UNSEAL_KEY_FILE}" ]; then
      echo "ERROR: Unseal key not found at ${UNSEAL_KEY_FILE}"
      exit 1
    fi

    UNSEAL_KEY=$(cat "${UNSEAL_KEY_FILE}")

    # Unseal with retry logic
    MAX_UNSEAL_RETRIES=3
    UNSEAL_RETRY=0
    UNSEAL_SUCCESS=false

    while [ ${UNSEAL_RETRY} -lt ${MAX_UNSEAL_RETRIES} ]; do
      UNSEAL_RESPONSE=$(wget -qO- --header="Content-Type: application/json" --post-data="{\"key\":\"${UNSEAL_KEY}\"}" "${VAULT_ADDR}/v1/sys/unseal" 2>&1)

      # Verify unseal was successful by checking sealed status
      sleep 1
      VERIFY_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/seal-status" 2>/dev/null || echo '{"sealed":true}')
      VERIFY_SEALED=$(echo "${VERIFY_STATUS}" | grep -o '"sealed":[^,}]*' | cut -d':' -f2)

      if [ "${VERIFY_SEALED}" = "false" ]; then
        UNSEAL_SUCCESS=true
        echo "OpenBao unsealed successfully"
        break
      fi

      UNSEAL_RETRY=$((UNSEAL_RETRY + 1))
      echo "Unseal attempt ${UNSEAL_RETRY} failed, retrying..."
      sleep 2
    done

    if [ "${UNSEAL_SUCCESS}" = "false" ]; then
      echo "ERROR: Failed to unseal OpenBao after ${MAX_UNSEAL_RETRIES} attempts"
      exit 1
    fi
  else
    echo "OpenBao is already unsealed"
  fi

  # Verify Transit engine and AppRole are configured
  if [ -f "${ROOT_TOKEN_FILE}" ]; then
    export VAULT_TOKEN=$(cat "${ROOT_TOKEN_FILE}")

    # Check Transit engine
    if ! bao secrets list | grep -q "transit/"; then
      echo "Transit engine not found - configuring..."
      bao secrets enable transit
      echo "Transit secrets engine enabled"
    else
      echo "Transit secrets engine already enabled"
    fi

    # Check AppRole
    if ! bao auth list | grep -q "approle/"; then
      echo "AppRole not found - configuring..."
      bao auth enable approle
      echo "AppRole auth method enabled"
    else
      echo "AppRole auth method already enabled"
    fi
  fi

  echo "Initialization check complete - OpenBao is ready"
  exit 0
fi

echo "OpenBao is not initialized - initializing with 1-of-1 key shares..."

# Initialize with 1 key share, threshold 1 (turnkey mode)
INIT_OUTPUT=$(bao operator init -key-shares=1 -key-threshold=1 -format=json)

# Extract unseal key and root token from JSON output
# First collapse multi-line JSON to single line, then parse
INIT_JSON=$(echo "${INIT_OUTPUT}" | tr -d '\n' | tr -d ' ')
UNSEAL_KEY=$(echo "${INIT_JSON}" | grep -o '"unseal_keys_b64":\["[^"]*"' | cut -d'"' -f4)
ROOT_TOKEN=$(echo "${INIT_JSON}" | grep -o '"root_token":"[^"]*"' | cut -d'"' -f4)

# Save to files
echo "${UNSEAL_KEY}" > "${UNSEAL_KEY_FILE}"
echo "${ROOT_TOKEN}" > "${ROOT_TOKEN_FILE}"
chmod 600 "${UNSEAL_KEY_FILE}" "${ROOT_TOKEN_FILE}"

echo "Initialization complete"
echo "Unseal key saved to ${UNSEAL_KEY_FILE}"
echo "Root token saved to ${ROOT_TOKEN_FILE}"

# Unseal with retry logic
echo "Unsealing OpenBao..."
MAX_UNSEAL_RETRIES=3
UNSEAL_RETRY=0
UNSEAL_SUCCESS=false

while [ ${UNSEAL_RETRY} -lt ${MAX_UNSEAL_RETRIES} ]; do
  UNSEAL_RESPONSE=$(wget -qO- --header="Content-Type: application/json" --post-data="{\"key\":\"${UNSEAL_KEY}\"}" "${VAULT_ADDR}/v1/sys/unseal" 2>&1)

  # Verify unseal was successful by checking sealed status
  sleep 1
  VERIFY_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/seal-status" 2>/dev/null || echo '{"sealed":true}')
  VERIFY_SEALED=$(echo "${VERIFY_STATUS}" | grep -o '"sealed":[^,}]*' | cut -d':' -f2)

  if [ "${VERIFY_SEALED}" = "false" ]; then
    UNSEAL_SUCCESS=true
    echo "OpenBao unsealed successfully"
    break
  fi

  UNSEAL_RETRY=$((UNSEAL_RETRY + 1))
  echo "Unseal attempt ${UNSEAL_RETRY} failed, retrying..."
  sleep 2
done

if [ "${UNSEAL_SUCCESS}" = "false" ]; then
  echo "ERROR: Failed to unseal OpenBao after ${MAX_UNSEAL_RETRIES} attempts"
  exit 1
fi

# Configure with root token
export VAULT_TOKEN="${ROOT_TOKEN}"

# Enable Transit secrets engine
echo "Enabling Transit secrets engine..."
bao secrets enable transit
echo "Transit secrets engine enabled"

# Create Transit encryption keys
echo "Creating Transit encryption keys..."

bao write -f transit/keys/mosaic-credentials type=aes256-gcm96
echo "Created key: mosaic-credentials"

bao write -f transit/keys/mosaic-account-tokens type=aes256-gcm96
echo "Created key: mosaic-account-tokens"

bao write -f transit/keys/mosaic-federation type=aes256-gcm96
echo "Created key: mosaic-federation"

bao write -f transit/keys/mosaic-llm-config type=aes256-gcm96
echo "Created key: mosaic-llm-config"

echo "All Transit keys created"

# Enable AppRole auth method
echo "Enabling AppRole auth method..."
bao auth enable approle
echo "AppRole auth method enabled"

# Create Transit-only policy
echo "Creating Transit-only policy..."
cat > /tmp/transit-policy.hcl <<EOF
# Transit-only policy for Mosaic API
# Allows encrypt/decrypt operations only

path "transit/encrypt/*" {
  capabilities = ["update"]
}

path "transit/decrypt/*" {
  capabilities = ["update"]
}

# Deny all other paths
path "*" {
  capabilities = ["deny"]
}
EOF

bao policy write mosaic-transit-policy /tmp/transit-policy.hcl
echo "Transit policy created"

# Create AppRole
echo "Creating AppRole: mosaic-transit..."
bao write auth/approle/role/mosaic-transit \
  token_policies="mosaic-transit-policy" \
  token_ttl=1h \
  token_max_ttl=4h
echo "AppRole created"

# Get AppRole credentials
echo "Generating AppRole credentials..."
ROLE_ID_JSON=$(bao read -format=json auth/approle/role/mosaic-transit/role-id | tr -d '\n' | tr -d ' ')
ROLE_ID=$(echo "${ROLE_ID_JSON}" | grep -o '"role_id":"[^"]*"' | cut -d'"' -f4)

SECRET_ID_JSON=$(bao write -format=json -f auth/approle/role/mosaic-transit/secret-id | tr -d '\n' | tr -d ' ')
SECRET_ID=$(echo "${SECRET_ID_JSON}" | grep -o '"secret_id":"[^"]*"' | cut -d'"' -f4)

# Save credentials to file
cat > "${APPROLE_CREDS_FILE}" <<EOF
{
  "role_id": "${ROLE_ID}",
  "secret_id": "${SECRET_ID}"
}
EOF
chmod 600 "${APPROLE_CREDS_FILE}"

echo "AppRole credentials saved to ${APPROLE_CREDS_FILE}"

echo "=== OpenBao Configuration Complete ==="
echo "Status:"
bao status

echo ""
echo "Transit Keys:"
bao list transit/keys

echo ""
echo "Auth Methods:"
bao auth list

echo ""
echo "Initialization complete - OpenBao is ready for use"

# Watch loop to handle unsealing after container restarts
echo ""
echo "Starting unseal watch loop (checks every 30 seconds)..."
while true; do
  sleep 30

  # Check if OpenBao is sealed
  SEAL_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/seal-status" 2>/dev/null || echo '{"sealed":false}')
  IS_SEALED=$(echo "${SEAL_STATUS}" | grep -o '"sealed":[^,}]*' | cut -d':' -f2)

  if [ "${IS_SEALED}" = "true" ]; then
    echo "OpenBao is sealed - unsealing..."
    if [ -f "${UNSEAL_KEY_FILE}" ]; then
      UNSEAL_KEY=$(cat "${UNSEAL_KEY_FILE}")

      # Try to unseal with verification
      UNSEAL_RESPONSE=$(wget -qO- --header="Content-Type: application/json" --post-data="{\"key\":\"${UNSEAL_KEY}\"}" "${VAULT_ADDR}/v1/sys/unseal" 2>&1)

      # Verify unseal was successful
      sleep 1
      VERIFY_STATUS=$(wget -qO- "${VAULT_ADDR}/v1/sys/seal-status" 2>/dev/null || echo '{"sealed":true}')
      VERIFY_SEALED=$(echo "${VERIFY_STATUS}" | grep -o '"sealed":[^,}]*' | cut -d':' -f2)

      if [ "${VERIFY_SEALED}" = "false" ]; then
        echo "OpenBao unsealed successfully"
      else
        echo "WARNING: Unseal operation completed but OpenBao is still sealed"
      fi
    else
      echo "WARNING: Unseal key not found, cannot auto-unseal"
    fi
  fi
done