#!/usr/bin/env bash
set -eo pipefail

# Colors for terminal output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# Helper function for logging
log() {
  local level=$1
  local message=$2
  local color=$NC

  case $level in
    "INFO") color=$GREEN ;;
    "WARN") color=$YELLOW ;;
    "ERROR") color=$RED ;;
  esac

  echo -e "${color}[$level] $message${NC}"
}

# Function to cleanup on exit
cleanup() {
  log "INFO" "Cleaning up resources..."
  docker-compose down -v > /dev/null 2>&1 || true
  rm -f vault-credentials.txt || true
  log "INFO" "Cleanup complete"
}

# Register the cleanup function to run on exit
trap cleanup EXIT

# Start the test
log "INFO" "Starting Docker Compose test for vault-hier"

# Check if Docker is running
if ! docker info > /dev/null 2>&1; then
  log "ERROR" "Docker is not running or not accessible. Please start Docker."
  exit 1
fi

# Check if docker-compose is available
if ! command -v docker-compose > /dev/null 2>&1; then
  log "ERROR" "docker-compose command not found. Please install Docker Compose."
  exit 1
fi

# Check if jq is available
if ! command -v jq > /dev/null 2>&1; then
  log "ERROR" "jq command not found. Please install jq (JSON processor)."
  exit 1
fi

# Build the Docker image
log "INFO" "Building Docker image..."
docker-compose build

# Clean up any existing containers and volumes
log "INFO" "Cleaning up any existing resources..."
docker-compose down -v

# Start fresh containers
log "INFO" "Starting containers..."
docker-compose up -d

# Function to wait for vault-init to complete
wait_for_vault_init() {
  local max_attempts=30
  local attempt=1
  local container_name="vault-init"
  
  while [ $attempt -le $max_attempts ]; do
    log "INFO" "Waiting for vault-init to complete (attempt $attempt/$max_attempts)..."
    
    # Check exit code of vault-init container
    status=$(docker inspect --format='{{.State.Status}}' $container_name 2>/dev/null || echo "error")
    
    if [ "$status" = "exited" ]; then
      exit_code=$(docker inspect --format='{{.State.ExitCode}}' $container_name)
      if [ "$exit_code" = "0" ]; then
        log "INFO" "vault-init completed successfully"
        return 0
      else
        log "ERROR" "vault-init failed with exit code $exit_code"
        docker logs $container_name
        return 1
      fi
    elif [ "$status" = "error" ]; then
      log "ERROR" "Container $container_name not found"
      return 1
    fi
    
    attempt=$((attempt + 1))
    sleep 2
  done
  
  log "ERROR" "Timed out waiting for vault-init to complete"
  docker logs $container_name
  return 1
}

# Wait for vault-init to complete
wait_for_vault_init

# Check if vault-credentials.json was created
if [ -f "vault-credentials.json" ]; then
  log "INFO" "JSON credentials file was created successfully"
else
  log "ERROR" "JSON credentials file was not created"
  exit 1
fi

# Verify the content of vault-credentials.json
if jq -e '.keys_base64 | length' vault-credentials.json >/dev/null && \
   jq -e '.root_token' vault-credentials.json >/dev/null; then
  log "INFO" "JSON credentials file contains expected content"
else
  log "ERROR" "JSON credentials file doesn't contain expected content"
  exit 1
fi

# Also check for backward compatibility
if [ -f "vault-credentials.txt" ]; then
  log "INFO" "Text credentials file was also created (for backward compatibility)"
fi

# Verify Vault is unsealed after initial setup
vault_status=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault status -format=json 2>/dev/null || echo '{"sealed": true}')

# Check if Vault is unsealed by looking for "sealed":false
sealed=$(echo "$vault_status" | grep -o '"sealed":[^,]*' || echo '"sealed":true')
log "INFO" "Seal status: $sealed"

if [[ "$sealed" == *"false"* ]]; then
  log "INFO" "Vault is properly unsealed after initial setup"
else
  log "ERROR" "Vault is still sealed after initial setup"
  echo $vault_status
  exit 1
fi

# Test restart scenario
log "INFO" "Testing restart scenario..."

# Stop and remove just the vault-init container
docker-compose rm -f -s vault-init

# Stop and restart vault
log "INFO" "Stopping Vault container..."
docker-compose stop vault

log "INFO" "Restarting Vault container..."
docker-compose start vault
sleep 5

# Verify Vault is sealed after restart (it should be)
vault_status=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault status -format=json 2>/dev/null || echo '{"sealed": true}')

sealed=$(echo "$vault_status" | grep -o '"sealed":[^,]*' || echo '"sealed":false')
log "INFO" "Seal status after restart: $sealed"

if [[ "$sealed" == *"true"* ]]; then
  log "INFO" "Vault is correctly sealed after restart"
else
  log "WARN" "Vault is not sealed after restart - this is unexpected"
  echo $vault_status
fi

# Extract keys and token from JSON credentials file
log "INFO" "Extracting unseal keys and root token from JSON credentials file..."
# Using jq to extract the token
root_token=$(jq -r '.root_token' vault-credentials.json)

# First, try running 'vault operator unseal' directly for a more robust test
log "INFO" "Attempting to unseal Vault directly with unseal keys..."
# Use jq to extract the keys directly into an array - more elegant
readarray -t unseal_keys_array < <(jq -r '.keys_base64[0:3][]' vault-credentials.json)

# Apply each key
for key in "${unseal_keys_array[@]}"; do
  if [ -n "$key" ]; then
    log "INFO" "Applying unseal key: ${key:0:8}..." # Show only first 8 chars for security
    docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault operator unseal "$key"
  fi
done

# As a fallback, also try running vault-init with environment variables
log "INFO" "Starting vault-init with environment variables..."
# Use the array to set environment variables
env_vars="-e VAULT_ADDR=http://vault:8200"
for i in "${!unseal_keys_array[@]}"; do
  env_vars="$env_vars -e VAULT_UNSEAL_KEY_$((i+1))=${unseal_keys_array[$i]}"
done

# Run the command with all environment variables
eval "docker-compose run $env_vars --rm vault-init"

# Verify Vault is unsealed now
vault_status=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault status -format=json 2>/dev/null || echo '{"sealed": true}')

sealed=$(echo "$vault_status" | grep -o '"sealed":[^,]*' || echo '"sealed":true')
log "INFO" "Seal status after unseal attempts: $sealed"

if [[ "$sealed" == *"false"* ]]; then
  log "INFO" "Vault was successfully unsealed after restart"
else
  log "ERROR" "Vault is still sealed after restart"
  echo $vault_status
  exit 1
fi

# Test some basic Vault operations
log "INFO" "Testing basic Vault operations..."

# Write a secret using the root token from JSON credentials
token_result=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault login "$root_token" 2>&1)
login_success=$(echo "$token_result" | grep -c "Success" || echo "0")
if [ "$login_success" -gt 0 ]; then
  log "INFO" "Successfully logged in with root token"
else
  log "ERROR" "Failed to log in with root token"
  echo "$token_result"
  exit 1
fi

# Enable KV secrets engine
enable_result=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault secrets enable -path=kv kv 2>&1 || echo "KV already enabled")
log "INFO" "Secrets engine: $enable_result"

# Write a test secret
write_result=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault kv put kv/test-secret username=vault-hier password=test123 2>&1)
log "INFO" "Write result: $(echo "$write_result" | grep "Success")"

# Read the test secret
read_result=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault kv get -format=json kv/test-secret 2>&1)
if echo "$read_result" | grep -q "vault-hier"; then
  log "INFO" "Successfully read back test secret"
else
  log "ERROR" "Failed to read back test secret"
  echo "$read_result"
  exit 1
fi

# All tests passed
log "INFO" "All tests passed successfully!"
exit 0