#!/bin/bash

# Foxhunt Vault Initialization Script
# This script initializes and unseals Vault, then sets up the basic configuration

set -e

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

# Script configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VAULT_ADDR="${VAULT_ADDR:-https://foxhunt-vault:8200}"
VAULT_INIT_FILE="${VAULT_INIT_FILE:-/vault-init/init.json}"
VAULT_TOKEN_FILE="${VAULT_TOKEN_FILE:-/vault-init/root_token}"
MAX_RETRIES=30
RETRY_DELAY=5

# Logging function
log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] ✓${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] ⚠${NC} $1"
}

log_error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ✗${NC} $1"
}

# Wait for Vault to be ready
wait_for_vault() {
    log "Waiting for Vault to be ready at ${VAULT_ADDR}..."
    
    local count=0
    while ! vault status > /dev/null 2>&1; do
        if [ $count -eq $MAX_RETRIES ]; then
            log_error "Vault did not become ready within expected time"
            return 1
        fi
        
        log "Vault not ready, waiting... (attempt $((count + 1))/${MAX_RETRIES})"
        sleep $RETRY_DELAY
        count=$((count + 1))
    done
    
    log_success "Vault is responding"
}

# Initialize Vault if not already initialized
initialize_vault() {
    log "Checking if Vault is already initialized..."
    
    if vault status | grep -q "Initialized.*true"; then
        log_warning "Vault is already initialized"
        
        if [ -f "$VAULT_INIT_FILE" ]; then
            log "Using existing initialization data"
            return 0
        else
            log_error "Vault is initialized but init file not found at $VAULT_INIT_FILE"
            log_error "Manual intervention required"
            return 1
        fi
    fi
    
    log "Initializing Vault..."
    
    # Create directory for init files
    mkdir -p "$(dirname "$VAULT_INIT_FILE")"
    
    # Initialize Vault with 5 key shares and threshold of 3
    vault operator init \
        -key-shares=5 \
        -key-threshold=3 \
        -format=json > "$VAULT_INIT_FILE"
    
    if [ $? -eq 0 ]; then
        log_success "Vault initialized successfully"
        
        # Extract and save root token
        jq -r '.root_token' "$VAULT_INIT_FILE" > "$VAULT_TOKEN_FILE"
        
        # Set secure permissions
        chmod 600 "$VAULT_INIT_FILE" "$VAULT_TOKEN_FILE"
        
        log "Initialization data saved to: $VAULT_INIT_FILE"
        log "Root token saved to: $VAULT_TOKEN_FILE"
        
        # Display unseal keys for manual storage
        log_warning "IMPORTANT: Store these unseal keys securely!"
        echo -e "${YELLOW}"
        jq -r '.unseal_keys_b64[]' "$VAULT_INIT_FILE" | nl -v0 -w2 -s': '
        echo -e "${NC}"
        
    else
        log_error "Failed to initialize Vault"
        return 1
    fi
}

# Unseal Vault
unseal_vault() {
    log "Checking Vault seal status..."
    
    if ! vault status | grep -q "Sealed.*true"; then
        log_success "Vault is already unsealed"
        return 0
    fi
    
    log "Unsealing Vault..."
    
    if [ ! -f "$VAULT_INIT_FILE" ]; then
        log_error "Initialization file not found: $VAULT_INIT_FILE"
        return 1
    fi
    
    # Extract unseal keys and unseal
    local unseal_keys=($(jq -r '.unseal_keys_b64[]' "$VAULT_INIT_FILE"))
    local threshold=$(jq -r '.secret_threshold' "$VAULT_INIT_FILE")
    
    log "Using threshold of $threshold unseal keys"
    
    for i in $(seq 0 $((threshold - 1))); do
        log "Providing unseal key $((i + 1)) of $threshold"
        echo "${unseal_keys[$i]}" | vault operator unseal -
        
        if [ $? -ne 0 ]; then
            log_error "Failed to provide unseal key $((i + 1))"
            return 1
        fi
    done
    
    # Verify unsealing
    if vault status | grep -q "Sealed.*false"; then
        log_success "Vault successfully unsealed"
    else
        log_error "Vault unsealing verification failed"
        return 1
    fi
}

# Authenticate with root token
authenticate() {
    log "Authenticating with Vault..."
    
    if [ ! -f "$VAULT_TOKEN_FILE" ]; then
        log_error "Root token file not found: $VAULT_TOKEN_FILE"
        return 1
    fi
    
    export VAULT_TOKEN=$(cat "$VAULT_TOKEN_FILE")
    
    # Verify authentication
    if vault auth -method=token "$VAULT_TOKEN" > /dev/null 2>&1; then
        log_success "Successfully authenticated with Vault"
    else
        log_error "Failed to authenticate with Vault"
        return 1
    fi
}

# Enable audit logging
enable_audit() {
    log "Enabling audit logging..."
    
    # Check if audit is already enabled
    if vault audit list | grep -q "file/"; then
        log_warning "Audit logging already enabled"
        return 0
    fi
    
    # Enable file audit device
    vault audit enable file file_path=/vault/logs/audit.log
    
    if [ $? -eq 0 ]; then
        log_success "Audit logging enabled"
    else
        log_warning "Failed to enable audit logging (may need manual configuration)"
    fi
}

# Enable KV v2 secrets engine
enable_kv_engine() {
    log "Enabling KV v2 secrets engine..."
    
    # Check if already enabled
    if vault secrets list | grep -q "foxhunt/"; then
        log_warning "KV secrets engine already enabled at foxhunt/"
        return 0
    fi
    
    # Enable KV v2 at foxhunt path
    vault secrets enable -path=foxhunt -version=2 kv
    
    if [ $? -eq 0 ]; then
        log_success "KV v2 secrets engine enabled at foxhunt/"
    else
        log_error "Failed to enable KV v2 secrets engine"
        return 1
    fi
}

# Enable AppRole authentication
enable_approle() {
    log "Enabling AppRole authentication method..."
    
    # Check if already enabled
    if vault auth list | grep -q "approle/"; then
        log_warning "AppRole authentication already enabled"
        return 0
    fi
    
    # Enable AppRole
    vault auth enable approle
    
    if [ $? -eq 0 ]; then
        log_success "AppRole authentication enabled"
    else
        log_error "Failed to enable AppRole authentication"
        return 1
    fi
}

# Load policies
load_policies() {
    log "Loading Vault policies..."
    
    local policies_dir="/vault/config/policies"
    
    if [ ! -d "$policies_dir" ]; then
        log_error "Policies directory not found: $policies_dir"
        return 1
    fi
    
    for policy_file in "$policies_dir"/*.hcl; do
        if [ -f "$policy_file" ]; then
            local policy_name=$(basename "$policy_file" .hcl)
            log "Loading policy: $policy_name"
            
            vault policy write "$policy_name" "$policy_file"
            
            if [ $? -eq 0 ]; then
                log_success "Policy '$policy_name' loaded successfully"
            else
                log_error "Failed to load policy '$policy_name'"
                return 1
            fi
        fi
    done
}

# Create service AppRoles
create_service_approles() {
    log "Creating service AppRoles..."
    
    # Trading Service AppRole
    log "Creating trading-service AppRole..."
    vault write auth/approle/role/trading-service \
        token_policies="trading-service" \
        token_ttl=1h \
        token_max_ttl=24h \
        bind_secret_id=true \
        secret_id_ttl=24h
    
    # Backtesting Service AppRole
    log "Creating backtesting-service AppRole..."
    vault write auth/approle/role/backtesting-service \
        token_policies="backtesting-service" \
        token_ttl=1h \
        token_max_ttl=24h \
        bind_secret_id=true \
        secret_id_ttl=24h
    
    # TLI Client AppRole
    log "Creating tli-client AppRole..."
    vault write auth/approle/role/tli-client \
        token_policies="tli-client" \
        token_ttl=30m \
        token_max_ttl=8h \
        bind_secret_id=true \
        secret_id_ttl=8h
    
    log_success "Service AppRoles created"
}

# Display service credentials
display_service_credentials() {
    log "Retrieving service credentials..."
    
    echo -e "\n${GREEN}=== SERVICE CREDENTIALS ===${NC}"
    echo -e "${YELLOW}Store these credentials securely for service configuration${NC}"
    
    # Trading Service
    echo -e "\n${BLUE}Trading Service:${NC}"
    echo -n "Role ID: "
    vault read -field=role_id auth/approle/role/trading-service/role-id
    echo -n "Secret ID: "
    vault write -field=secret_id -f auth/approle/role/trading-service/secret-id
    
    # Backtesting Service
    echo -e "\n${BLUE}Backtesting Service:${NC}"
    echo -n "Role ID: "
    vault read -field=role_id auth/approle/role/backtesting-service/role-id
    echo -n "Secret ID: "
    vault write -field=secret_id -f auth/approle/role/backtesting-service/secret-id
    
    # TLI Client
    echo -e "\n${BLUE}TLI Client:${NC}"
    echo -n "Role ID: "
    vault read -field=role_id auth/approle/role/tli-client/role-id
    echo -n "Secret ID: "
    vault write -field=secret_id -f auth/approle/role/tli-client/secret-id
    
    echo -e "\n${GREEN}=== INITIALIZATION COMPLETE ===${NC}"
}

# Main execution
main() {
    log "Starting Foxhunt Vault initialization..."
    
    # Wait for Vault to be ready
    wait_for_vault || exit 1
    
    # Initialize Vault
    initialize_vault || exit 1
    
    # Unseal Vault
    unseal_vault || exit 1
    
    # Authenticate
    authenticate || exit 1
    
    # Enable audit logging
    enable_audit
    
    # Enable KV secrets engine
    enable_kv_engine || exit 1
    
    # Enable AppRole authentication
    enable_approle || exit 1
    
    # Load policies
    load_policies || exit 1
    
    # Create service AppRoles
    create_service_approles || exit 1
    
    # Display credentials
    display_service_credentials
    
    log_success "Vault initialization completed successfully!"
    log "Next steps:"
    log "1. Store the unseal keys and root token securely"
    log "2. Configure services with their AppRole credentials"
    log "3. Populate secrets using the appropriate setup script"
}

# Handle signals
trap 'log_error "Script interrupted"; exit 1' INT TERM

# Set Vault address
export VAULT_ADDR

# Run main function
main "$@"