#!/bin/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

# 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.txt was created
if [ -f "vault-credentials.txt" ]; then
  log "INFO" "Credentials file was created successfully"
else
  log "ERROR" "Credentials file was not created"
  exit 1
fi

# Verify the content of vault-credentials.txt
if grep -q "Unseal Keys:" vault-credentials.txt && grep -q "Root Token:" vault-credentials.txt; then
  log "INFO" "Credentials file contains expected content"
else
  log "ERROR" "Credentials file doesn't contain expected content"
  exit 1
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
if echo "$vault_status" | grep -q '"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}')

if echo "$vault_status" | grep -q '"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 from credentials file and root token
log "INFO" "Extracting unseal keys and root token from credentials file..."
unseal_keys=$(grep "Base64 Unseal Keys:" -A 3 vault-credentials.txt | grep "Key" | awk '{print $3}')
root_token=$(grep "Root Token:" vault-credentials.txt | awk '{print $3}')

# First, try running 'vault operator unseal' directly for a more robust test
log "INFO" "Attempting to unseal Vault directly with unseal keys..."
key1=$(echo "$unseal_keys" | head -n 1)
key2=$(echo "$unseal_keys" | head -n 2 | tail -n 1)
key3=$(echo "$unseal_keys" | head -n 3 | tail -n 1)

docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault operator unseal "$key1"
docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault operator unseal "$key2"
docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault operator unseal "$key3"

# As a fallback, also try running vault-init with environment variables
log "INFO" "Starting vault-init with environment variables..."
docker-compose run -e VAULT_ADDR=http://vault:8200 \
  -e VAULT_UNSEAL_KEY_1=$(echo "$unseal_keys" | head -n 1) \
  -e VAULT_UNSEAL_KEY_2=$(echo "$unseal_keys" | head -n 2 | tail -n 1) \
  -e VAULT_UNSEAL_KEY_3=$(echo "$unseal_keys" | head -n 3 | tail -n 1) \
  --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}')

if echo "$vault_status" | grep -q '"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
token_result=$(docker-compose exec -T vault env VAULT_ADDR=http://127.0.0.1:8200 vault login "$root_token" 2>&1)
log "INFO" "Login result: $(echo "$token_result" | grep "Success")"

# 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