#!/bin/sh
# Copyright (c) Abstract Machines
# SPDX-License-Identifier: Apache-2.0

set -e

apk add --no-cache jq

# Create required directories
mkdir -p /opt/openbao/config /opt/openbao/data /opt/openbao/logs

cat > /opt/openbao/config/config.hcl << 'EOF'
storage "file" {
  path = "/opt/openbao/data"
}
listener "tcp" {
  address     = "0.0.0.0:8200"
  tls_disable = true
}
ui = true
log_level = "Info"
disable_mlock = true
# API timeout settings
default_lease_ttl = "168h"
max_lease_ttl = "720h"
EOF

export BAO_ADDR=http://127.0.0.1:8200

wait_for_bao() {
  local retries=30
  while [ "$retries" -gt 0 ]; do
    local code=0
    bao status -format=json >/dev/null 2>&1 || code=$?
    # exit 0 = unsealed, exit 2 = sealed — both mean the server is up and responding
    if [ "$code" -eq 0 ] || [ "$code" -eq 2 ]; then
      return 0
    fi
    retries=$((retries - 1))
    sleep 2
  done
  echo "ERROR: OpenBao server did not become ready in time" >&2
  exit 1
}

unseal_bao() {
  local key1="$1" key2="$2" key3="$3"
  if [ -z "$key1" ] || [ -z "$key2" ] || [ -z "$key3" ]; then
    echo "ERROR: One or more unseal keys are empty — cannot unseal" >&2
    exit 1
  fi
  bao operator unseal "$key1"
  bao operator unseal "$key2"
  bao operator unseal "$key3"
}

create_pki_policy() {
  cat > /opt/openbao/config/pki-policy.hcl << EOF
path "pki_int/issue/${MG_CERTS_OPENBAO_PKI_ROLE}" {
  capabilities = ["create", "update"]
}
path "pki_int/sign/${MG_CERTS_OPENBAO_PKI_ROLE}" {
  capabilities = ["create", "update"]
}
path "pki_int/sign-verbatim/${MG_CERTS_OPENBAO_PKI_ROLE}" {
  capabilities = ["create", "update"]
}
path "pki_int/certs" {
  capabilities = ["list"]
}
path "pki_int/cert/*" {
  capabilities = ["read"]
}
path "pki_int/revoke" {
  capabilities = ["create", "update"]
}
path "pki_int/ca" {
  capabilities = ["read"]
}
path "pki_int/ca_chain" {
  capabilities = ["read"]
}
path "pki_int/crl" {
  capabilities = ["read"]
}
path "pki/ca" {
  capabilities = ["read"]
}
path "pki/ca_chain" {
  capabilities = ["read"]
}
path "pki/crl" {
  capabilities = ["read"]
}
path "auth/token/renew-self" {
  capabilities = ["update"]
}
path "auth/token/lookup-self" {
  capabilities = ["read"]
}
path "sys/renew/*" {
  capabilities = ["update"]
}
path "auth/approle/role/${MG_CERTS_OPENBAO_PKI_ROLE}/secret-id" {
  capabilities = ["create", "update"]
}
path "auth/approle/role/${MG_CERTS_OPENBAO_PKI_ROLE}/secret-id-accessor/lookup" {
  capabilities = ["create", "update"]
}
path "auth/approle/role/${MG_CERTS_OPENBAO_PKI_ROLE}/secret-id-accessor/destroy" {
  capabilities = ["create", "update"]
}
EOF
  bao policy write pki-policy /opt/openbao/config/pki-policy.hcl > /dev/null
}

# Check if we have pre-configured unseal keys and root token
if [ -n "$MG_CERTS_OPENBAO_UNSEAL_KEY_1" ] && [ -n "$MG_CERTS_OPENBAO_UNSEAL_KEY_2" ] && [ -n "$MG_CERTS_OPENBAO_UNSEAL_KEY_3" ] && [ -n "$MG_CERTS_OPENBAO_ROOT_TOKEN" ]; then
  echo "Using pre-configured unseal keys and root token..."
  bao server -config=/opt/openbao/config/config.hcl > /opt/openbao/logs/server.log 2>&1 &
  BAO_PID=$!
  wait_for_bao

  unseal_bao "$MG_CERTS_OPENBAO_UNSEAL_KEY_1" "$MG_CERTS_OPENBAO_UNSEAL_KEY_2" "$MG_CERTS_OPENBAO_UNSEAL_KEY_3"

  export BAO_TOKEN=$MG_CERTS_OPENBAO_ROOT_TOKEN
else
  # Initialize OpenBao if not already done
  if [ ! -f /opt/openbao/data/init.json ]; then
    echo "Initializing OpenBao for the first time..."
    bao server -config=/opt/openbao/config/config.hcl > /opt/openbao/logs/server.log 2>&1 &
    BAO_PID=$!
    wait_for_bao

    # Initialize with 5 key shares and threshold of 3
    bao operator init -key-shares=5 -key-threshold=3 -format=json > /opt/openbao/data/init.json

    # Extract unseal keys and root token
    UNSEAL_KEY_1=$(jq -r '.unseal_keys_b64[0]' /opt/openbao/data/init.json)
    UNSEAL_KEY_2=$(jq -r '.unseal_keys_b64[1]' /opt/openbao/data/init.json)
    UNSEAL_KEY_3=$(jq -r '.unseal_keys_b64[2]' /opt/openbao/data/init.json)
    ROOT_TOKEN=$(jq -r '.root_token' /opt/openbao/data/init.json)

    unseal_bao "$UNSEAL_KEY_1" "$UNSEAL_KEY_2" "$UNSEAL_KEY_3"

    export BAO_TOKEN=$ROOT_TOKEN
    echo "OpenBao initialized successfully!"
  else
    echo "OpenBao already initialized, starting server..."
    bao server -config=/opt/openbao/config/config.hcl > /opt/openbao/logs/server.log 2>&1 &
    BAO_PID=$!
    wait_for_bao

    # Check if OpenBao is sealed and unseal if necessary
    if bao status -format=json | jq -e '.sealed == true' >/dev/null 2>&1; then
      echo "OpenBao is sealed, unsealing..."
      UNSEAL_KEY_1=$(jq -r '.unseal_keys_b64[0]' /opt/openbao/data/init.json)
      UNSEAL_KEY_2=$(jq -r '.unseal_keys_b64[1]' /opt/openbao/data/init.json)
      UNSEAL_KEY_3=$(jq -r '.unseal_keys_b64[2]' /opt/openbao/data/init.json)

      unseal_bao "$UNSEAL_KEY_1" "$UNSEAL_KEY_2" "$UNSEAL_KEY_3"
      echo "OpenBao unsealed successfully!"
    else
      echo "OpenBao is already unsealed!"
    fi

    ROOT_TOKEN=$(jq -r '.root_token' /opt/openbao/data/init.json)
    export BAO_TOKEN=$ROOT_TOKEN
  fi
fi

# Configure OpenBao PKI and AppRole if not already configured
if [ ! -f /opt/openbao/data/configured ]; then
  echo "Configuring OpenBao PKI and AppRole..."
  
  # Create namespace if specified
  if [ -n "$MG_CERTS_OPENBAO_NAMESPACE" ]; then
    if bao namespace create "$MG_CERTS_OPENBAO_NAMESPACE" 2>/tmp/ns_error; then
      export BAO_NAMESPACE="$MG_CERTS_OPENBAO_NAMESPACE"
      echo "$MG_CERTS_OPENBAO_NAMESPACE" > /opt/openbao/data/namespace
      echo "Created namespace: $MG_CERTS_OPENBAO_NAMESPACE"
    else
      if grep -q "namespace already exists" /tmp/ns_error; then
        export BAO_NAMESPACE="$MG_CERTS_OPENBAO_NAMESPACE"
        echo "$MG_CERTS_OPENBAO_NAMESPACE" > /opt/openbao/data/namespace
        echo "Using existing namespace: $MG_CERTS_OPENBAO_NAMESPACE"
      else
        echo "ERROR: Failed to create namespace $MG_CERTS_OPENBAO_NAMESPACE:" >&2
        cat /tmp/ns_error >&2
        exit 1
      fi
    fi
    rm -f /tmp/ns_error
  fi

  # Enable authentication methods and secrets engines
  if ! bao auth enable approle > /tmp/auth_success 2>/tmp/auth_error; then
    if ! grep -q "already in use" /tmp/auth_error; then
      echo "ERROR: Failed to enable AppRole auth method:" >&2
      cat /tmp/auth_error >&2
      exit 1
    fi
    echo "AppRole already enabled"
  fi
  rm -f /tmp/auth_error /tmp/auth_success

  # Enable PKI secrets engine
  if ! bao secrets enable -path=pki pki > /tmp/pki_success 2>/tmp/pki_error; then
    # If the failure wasn’t because the mount already exists, abort
    if ! grep -q "already in use" /tmp/pki_error; then
      echo "ERROR: Failed to enable PKI secrets engine:" >&2
      cat /tmp/pki_error >&2
      exit 1
    fi
    echo "PKI already enabled"
  fi
  rm -f /tmp/pki_error /tmp/pki_success

  # Configure PKI engine
  bao secrets tune -max-lease-ttl=87600h pki > /dev/null

  # Validate required CA environment variables
  for var in MG_CERTS_OPENBAO_PKI_CA_CN MG_CERTS_OPENBAO_PKI_CA_O MG_CERTS_OPENBAO_PKI_CA_C; do
    eval "value=\$var"
    if [ -z "$value" ]; then
      echo "ERROR: Required environment variable $var is not set" >&2
      exit 1
    fi
  done

  PKI_CMD="bao write -field=certificate pki/root/generate/internal \
    common_name=\"$MG_CERTS_OPENBAO_PKI_CA_CN\" \
    organization=\"$MG_CERTS_OPENBAO_PKI_CA_O\" \
    country=\"$MG_CERTS_OPENBAO_PKI_CA_C\" \
    ttl=87600h \
    key_bits=2048 \
    exclude_cn_from_sans=false"

  [ -n "$MG_CERTS_OPENBAO_PKI_CA_OU" ] && PKI_CMD="$PKI_CMD ou=\"$MG_CERTS_OPENBAO_PKI_CA_OU\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_L" ] && PKI_CMD="$PKI_CMD locality=\"$MG_CERTS_OPENBAO_PKI_CA_L\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_ST" ] && PKI_CMD="$PKI_CMD province=\"$MG_CERTS_OPENBAO_PKI_CA_ST\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_ADDR" ] && PKI_CMD="$PKI_CMD street_address=\"$MG_CERTS_OPENBAO_PKI_CA_ADDR\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_PO" ] && PKI_CMD="$PKI_CMD postal_code=\"$MG_CERTS_OPENBAO_PKI_CA_PO\""
  
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_DNS_NAMES" ] && PKI_CMD="$PKI_CMD alt_names=\"$MG_CERTS_OPENBAO_PKI_CA_DNS_NAMES\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_IP_ADDRESSES" ] && PKI_CMD="$PKI_CMD ip_sans=\"$MG_CERTS_OPENBAO_PKI_CA_IP_ADDRESSES\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_URI_SANS" ] && PKI_CMD="$PKI_CMD uri_sans=\"$MG_CERTS_OPENBAO_PKI_CA_URI_SANS\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_EMAIL_ADDRESSES" ] && PKI_CMD="$PKI_CMD email_sans=\"$MG_CERTS_OPENBAO_PKI_CA_EMAIL_ADDRESSES\""

  eval $PKI_CMD > /dev/null

  if [ $? -eq 0 ]; then
    echo "OpenBao root CA certificate generated successfully!"
  else
    echo "ERROR: Failed to generate OpenBao root CA certificate" >&2
    exit 1
  fi

  if ! bao secrets enable -path=pki_int pki > /tmp/pki_int_success 2>/tmp/pki_int_error; then
    if ! grep -q "already in use" /tmp/pki_int_error; then
      echo "ERROR: Failed to enable intermediate PKI secrets engine:" >&2
      cat /tmp/pki_int_error >&2
      exit 1
    fi
    echo "Intermediate PKI already enabled"
  fi
  rm -f /tmp/pki_int_error /tmp/pki_int_success

  bao secrets tune -max-lease-ttl=8760h pki_int > /dev/null

  INTERMEDIATE_CN="${MG_CERTS_OPENBAO_PKI_CA_CN} Intermediate"
  INTERMEDIATE_CSR_CMD="bao write -field=csr pki_int/intermediate/generate/internal \
    common_name=\"$INTERMEDIATE_CN\" \
    organization=\"$MG_CERTS_OPENBAO_PKI_CA_O\" \
    country=\"$MG_CERTS_OPENBAO_PKI_CA_C\" \
    ttl=8760h \
    key_bits=2048"

  [ -n "$MG_CERTS_OPENBAO_PKI_CA_OU" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD ou=\"$MG_CERTS_OPENBAO_PKI_CA_OU\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_L" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD locality=\"$MG_CERTS_OPENBAO_PKI_CA_L\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_ST" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD province=\"$MG_CERTS_OPENBAO_PKI_CA_ST\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_ADDR" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD street_address=\"$MG_CERTS_OPENBAO_PKI_CA_ADDR\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_PO" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD postal_code=\"$MG_CERTS_OPENBAO_PKI_CA_PO\""
  
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_DNS_NAMES" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD alt_names=\"$MG_CERTS_OPENBAO_PKI_CA_DNS_NAMES\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_IP_ADDRESSES" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD ip_sans=\"$MG_CERTS_OPENBAO_PKI_CA_IP_ADDRESSES\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_URI_SANS" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD uri_sans=\"$MG_CERTS_OPENBAO_PKI_CA_URI_SANS\""
  [ -n "$MG_CERTS_OPENBAO_PKI_CA_EMAIL_ADDRESSES" ] && INTERMEDIATE_CSR_CMD="$INTERMEDIATE_CSR_CMD email_sans=\"$MG_CERTS_OPENBAO_PKI_CA_EMAIL_ADDRESSES\""

  INTERMEDIATE_CSR=$(eval $INTERMEDIATE_CSR_CMD)

  if [ $? -ne 0 ] || [ -z "$INTERMEDIATE_CSR" ]; then
    echo "ERROR: Failed to generate intermediate CA CSR" >&2
    exit 1
  fi

  echo "Intermediate CA CSR generated successfully!"

  INTERMEDIATE_CERT=$(bao write -field=certificate pki/root/sign-intermediate \
    csr="$INTERMEDIATE_CSR" \
    format=pem_bundle \
    ttl=8760h \
    use_csr_values=true)

  if [ $? -ne 0 ] || [ -z "$INTERMEDIATE_CERT" ]; then
    echo "ERROR: Failed to sign intermediate CA certificate" >&2
    exit 1
  fi

  echo "Intermediate CA certificate signed successfully!"

  bao write pki/config/urls \
    issuing_certificates='http://127.0.0.1:8200/v1/pki/ca' \
    crl_distribution_points='http://127.0.0.1:8200/v1/pki/crl' \
    ocsp_servers='http://127.0.0.1:8200/v1/pki/ocsp' > /dev/null

  bao write pki_int/config/urls \
    issuing_certificates='http://127.0.0.1:8200/v1/pki_int/ca' \
    crl_distribution_points='http://127.0.0.1:8200/v1/pki_int/crl' \
    ocsp_servers='http://127.0.0.1:8200/v1/pki_int/ocsp' > /dev/null

  bao write pki_int/intermediate/set-signed certificate="$INTERMEDIATE_CERT" > /dev/null

  if [ $? -eq 0 ]; then
    echo "Intermediate CA setup completed successfully!"
  else
    echo "ERROR: Failed to set signed intermediate certificate" >&2
    exit 1
  fi

  echo "$INTERMEDIATE_CERT" > /opt/openbao/data/intermediate_ca.pem

  ROLE_CMD="bao write pki_int/roles/${MG_CERTS_OPENBAO_PKI_ROLE} \
    allow_any_name=true \
    enforce_hostnames=false \
    allow_ip_sans=true \
    allow_localhost=true \
    allow_bare_domains=true \
    allow_subdomains=true \
    allow_glob_domains=true \
    allowed_domains=\"*\" \
    allowed_uri_sans=\"*\" \
    allowed_other_sans=\"*\" \
    server_flag=true \
    client_flag=true \
    code_signing_flag=false \
    email_protection_flag=false \
    key_type=rsa \
    key_bits=2048 \
    key_usage=\"DigitalSignature,KeyEncipherment,KeyAgreement\" \
    ext_key_usage=\"ServerAuth,ClientAuth,OCSPSigning\" \
    use_csr_common_name=true \
    use_csr_sans=true \
    basic_constraints_valid_for_non_ca=true \
    max_ttl=720h \
    ttl=720h"

  eval "$ROLE_CMD" > /dev/null

  create_pki_policy

  # Create AppRole
  SECRET_ID_TTL="${MG_CERTS_OPENBAO_SECRET_ID_TTL}"
  bao write auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}" \
    token_policies=pki-policy \
    token_ttl=1h \
    token_max_ttl=4h \
    bind_secret_id=true \
    secret_id_ttl="$SECRET_ID_TTL" > /dev/null

  # Set custom role ID if provided
  if [ -n "$MG_CERTS_OPENBAO_APP_ROLE" ]; then
    bao write auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/role-id role_id="$MG_CERTS_OPENBAO_APP_ROLE" > /dev/null
  fi

  # Set custom secret ID if provided, otherwise generate one
  if [ -n "$MG_CERTS_OPENBAO_APP_SECRET" ]; then
    bao write auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/custom-secret-id secret_id="$MG_CERTS_OPENBAO_APP_SECRET" > /dev/null
    echo "$MG_CERTS_OPENBAO_APP_SECRET" > /opt/openbao/data/secret_id
  else
    GENERATED_SECRET_ID=$(bao write -field=secret_id -force auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/secret-id)
    echo "$GENERATED_SECRET_ID" > /opt/openbao/data/secret_id
  fi

  # Generate service token for additional access
  SERVICE_TOKEN=$(bao write -field=token auth/token/create \
    policies=pki-policy \
    ttl=24h \
    renewable=true \
    display_name="certs-service" 2>/dev/null)

  echo "SERVICE_TOKEN=$SERVICE_TOKEN" > /opt/openbao/data/service_token
  
  # Mark configuration as complete
  touch /opt/openbao/data/configured
  echo "OpenBao configuration completed successfully!"
else
  echo "OpenBao already configured, verifying and updating configuration..."
  
  # Restore namespace if it exists
  if [ -f /opt/openbao/data/namespace ] && [ -n "$MG_CERTS_OPENBAO_NAMESPACE" ]; then
    SAVED_NAMESPACE=$(cat /opt/openbao/data/namespace)
    if [ "$SAVED_NAMESPACE" = "$MG_CERTS_OPENBAO_NAMESPACE" ]; then
      export BAO_NAMESPACE="$MG_CERTS_OPENBAO_NAMESPACE"
    fi
  fi
  
  # Check if AppRole role exists, create if missing
  if ! bao read auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}" > /dev/null 2>&1; then
    if ! bao auth enable approle > /tmp/auth_success 2>/tmp/auth_error; then
      if ! grep -q "already in use" /tmp/auth_error; then
        echo "ERROR: Failed to enable AppRole auth method:" >&2
        cat /tmp/auth_error >&2
        exit 1
      fi
    fi
    rm -f /tmp/auth_error /tmp/auth_success
    
    create_pki_policy
    
    SECRET_ID_TTL="${MG_CERTS_OPENBAO_SECRET_ID_TTL}"
    bao write auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}" \
      token_policies=pki-policy \
      token_ttl=1h \
      token_max_ttl=4h \
      bind_secret_id=true \
      secret_id_ttl="$SECRET_ID_TTL" > /dev/null
    
    if [ -n "$MG_CERTS_OPENBAO_APP_ROLE" ]; then
      bao write auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/role-id role_id="$MG_CERTS_OPENBAO_APP_ROLE" > /dev/null
    fi
  fi
  
  SECRET_ID_VALID=false
  if [ -n "$MG_CERTS_OPENBAO_APP_SECRET" ]; then
    if bao write -field=client_token auth/approle/login role_id="$MG_CERTS_OPENBAO_APP_ROLE" secret_id="$MG_CERTS_OPENBAO_APP_SECRET" > /dev/null 2>&1; then
      SECRET_ID_VALID=true
      echo "$MG_CERTS_OPENBAO_APP_SECRET" > /opt/openbao/data/secret_id
    fi
  elif [ -f /opt/openbao/data/secret_id ]; then
    STORED_SECRET_ID=$(cat /opt/openbao/data/secret_id)
    if [ -n "$STORED_SECRET_ID" ]; then
      ROLE_ID=$(bao read -field=role_id auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/role-id)
      if bao write -field=client_token auth/approle/login role_id="$ROLE_ID" secret_id="$STORED_SECRET_ID" > /dev/null 2>&1; then
        SECRET_ID_VALID=true
      fi
    fi
  fi
  
  if [ "$SECRET_ID_VALID" = "false" ]; then
    NEW_SECRET_ID=$(bao write -field=secret_id -force auth/approle/role/"${MG_CERTS_OPENBAO_PKI_ROLE}"/secret-id)
    
    if [ -z "$NEW_SECRET_ID" ]; then
      echo "ERROR: Failed to generate new secret ID" >&2
    else
      echo "$NEW_SECRET_ID" > /opt/openbao/data/secret_id
      echo "Generated new secret ID for certs service"
    fi
  fi
  
  # Regenerate service token
  SERVICE_TOKEN=$(bao write -field=token auth/token/create \
    policies=pki-policy \
    ttl=24h \
    renewable=true \
    display_name="certs-service" 2>/dev/null)
  
  if [ -n "$SERVICE_TOKEN" ]; then
    echo "SERVICE_TOKEN=$SERVICE_TOKEN" > /opt/openbao/data/service_token
  fi
fi

echo "================================"
echo "OpenBao Production Setup Complete"
echo "================================"
echo "OpenBao Address: http://localhost:8200"
echo "UI Available at: http://localhost:8200/ui"
echo "================================"
echo "IMPORTANT: Store the init.json file securely!"
echo "It contains unseal keys and root token!"
echo "================================"

echo "OpenBao is ready for certs service on port 8200"

if [ -n "$BAO_PID" ]; then
  wait $BAO_PID
else
  echo "ERROR: OpenBao server process ID not available" >&2
  exit 1
fi