#!/bin/bash
#
# TLS Setup Validation Script for Foxhunt HFT Trading System
#
# This script validates the TLS configuration and certificate setup:
# - Verifies tonic TLS features are enabled
# - Checks certificate directory structure
# - Validates HashiCorp Vault connectivity (optional)
# - Tests TLS configuration loading
# - Runs integration tests

set -euo pipefail

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

# Configuration
FOXHUNT_ROOT="/home/jgrusewski/Work/foxhunt"
TLS_CERT_DIR="/opt/foxhunt/tls"
VAULT_CERT_CACHE="/opt/foxhunt/certs"

echo -e "${BLUE}🔒 Foxhunt TLS Setup Validation${NC}"
echo "================================================"

# Function to print status
print_status() {
    local status=$1
    local message=$2
    if [ "$status" = "OK" ]; then
        echo -e "${GREEN}✅ $message${NC}"
    elif [ "$status" = "WARN" ]; then
        echo -e "${YELLOW}⚠️  $message${NC}"
    else
        echo -e "${RED}❌ $message${NC}"
    fi
}

# Check if running as root for certificate directory creation
check_permissions() {
    echo -e "\n${BLUE}Checking Permissions...${NC}"
    
    if [ "$EUID" -eq 0 ]; then
        print_status "WARN" "Running as root - not recommended for production"
    else
        print_status "OK" "Running as non-root user"
    fi
}

# Verify tonic TLS features are enabled in Cargo.toml files
check_tonic_features() {
    echo -e "\n${BLUE}Checking Tonic TLS Features...${NC}"
    
    local files=(
        "$FOXHUNT_ROOT/tli/Cargo.toml"
        "$FOXHUNT_ROOT/services/trading_service/Cargo.toml"
        "$FOXHUNT_ROOT/services/backtesting_service/Cargo.toml"
        "$FOXHUNT_ROOT/services/ml_training_service/Cargo.toml"
    )
    
    local all_good=true
    
    for file in "${files[@]}"; do
        if [ -f "$file" ]; then
            if grep -q 'features.*=.*\["tls"' "$file"; then
                print_status "OK" "TLS features enabled in $(basename "$(dirname "$file")")"
            else
                print_status "FAIL" "TLS features NOT enabled in $(basename "$(dirname "$file")")"
                all_good=false
            fi
        else
            print_status "WARN" "File not found: $file"
        fi
    done
    
    if [ "$all_good" = true ]; then
        print_status "OK" "All Cargo.toml files have TLS features enabled"
    else
        print_status "FAIL" "Some Cargo.toml files missing TLS features"
        exit 1
    fi
}

# Check certificate directories and create if needed
check_certificate_directories() {
    echo -e "\n${BLUE}Checking Certificate Directories...${NC}"
    
    local dirs=(
        "$TLS_CERT_DIR"
        "$VAULT_CERT_CACHE"
    )
    
    for dir in "${dirs[@]}"; do
        if [ -d "$dir" ]; then
            print_status "OK" "Directory exists: $dir"
            
            # Check permissions
            if [ -r "$dir" ] && [ -w "$dir" ]; then
                print_status "OK" "Directory permissions OK: $dir"
            else
                print_status "WARN" "Directory permissions may be restrictive: $dir"
            fi
        else
            print_status "WARN" "Creating directory: $dir"
            sudo mkdir -p "$dir"
            sudo chown $USER:$USER "$dir"
            sudo chmod 755 "$dir"
        fi
    done
}

# Generate test certificates for development
generate_test_certificates() {
    echo -e "\n${BLUE}Generating Test Certificates...${NC}"
    
    local ca_key="$TLS_CERT_DIR/ca.key"
    local ca_crt="$TLS_CERT_DIR/ca.crt"
    local server_key="$TLS_CERT_DIR/server.key"
    local server_crt="$TLS_CERT_DIR/server.crt"
    local client_key="$TLS_CERT_DIR/client.key"
    local client_crt="$TLS_CERT_DIR/client.crt"
    
    # Check if certificates already exist
    if [ -f "$ca_crt" ] && [ -f "$server_crt" ] && [ -f "$client_crt" ]; then
        print_status "OK" "Test certificates already exist"
        return 0
    fi
    
    # Check if OpenSSL is available
    if ! command -v openssl &> /dev/null; then
        print_status "FAIL" "OpenSSL not found - cannot generate test certificates"
        return 1
    fi
    
    print_status "WARN" "Generating test certificates (DO NOT USE IN PRODUCTION)"
    
    # Generate CA private key
    openssl genpkey -algorithm RSA -out "$ca_key" -pkcs8 -pass pass:foxhunt
    
    # Generate CA certificate
    openssl req -new -x509 -key "$ca_key" -out "$ca_crt" -days 365 \
        -passin pass:foxhunt \
        -subj "/C=US/ST=CA/L=San Francisco/O=Foxhunt Trading/OU=Test CA/CN=Foxhunt Test CA"
    
    # Generate server private key
    openssl genpkey -algorithm RSA -out "$server_key" -pkcs8
    
    # Generate server certificate signing request
    openssl req -new -key "$server_key" -out "$TLS_CERT_DIR/server.csr" \
        -subj "/C=US/ST=CA/L=San Francisco/O=Foxhunt Trading/OU=Trading/CN=trading.foxhunt.internal"
    
    # Generate server certificate
    openssl x509 -req -in "$TLS_CERT_DIR/server.csr" -CA "$ca_crt" -CAkey "$ca_key" \
        -CAcreateserial -out "$server_crt" -days 365 -passin pass:foxhunt \
        -extensions v3_req -extfile <(echo "
[v3_req]
subjectAltName = @alt_names

[alt_names]
DNS.1 = trading.foxhunt.internal
DNS.2 = localhost
IP.1 = 127.0.0.1
")
    
    # Generate client private key
    openssl genpkey -algorithm RSA -out "$client_key" -pkcs8
    
    # Generate client certificate signing request
    openssl req -new -key "$client_key" -out "$TLS_CERT_DIR/client.csr" \
        -subj "/C=US/ST=CA/L=San Francisco/O=Foxhunt Trading/OU=Client/CN=client.foxhunt.internal"
    
    # Generate client certificate
    openssl x509 -req -in "$TLS_CERT_DIR/client.csr" -CA "$ca_crt" -CAkey "$ca_key" \
        -CAcreateserial -out "$client_crt" -days 365 -passin pass:foxhunt
    
    # Set proper permissions
    chmod 600 "$ca_key" "$server_key" "$client_key"
    chmod 644 "$ca_crt" "$server_crt" "$client_crt"
    
    # Clean up CSR files
    rm -f "$TLS_CERT_DIR/server.csr" "$TLS_CERT_DIR/client.csr"
    
    print_status "OK" "Test certificates generated successfully"
}

# Validate certificate files
validate_certificates() {
    echo -e "\n${BLUE}Validating Certificates...${NC}"
    
    local ca_crt="$TLS_CERT_DIR/ca.crt"
    local server_crt="$TLS_CERT_DIR/server.crt"
    local client_crt="$TLS_CERT_DIR/client.crt"
    
    if [ ! -f "$ca_crt" ] || [ ! -f "$server_crt" ] || [ ! -f "$client_crt" ]; then
        print_status "FAIL" "Certificate files missing"
        return 1
    fi
    
    # Validate CA certificate
    if openssl x509 -in "$ca_crt" -noout -text &>/dev/null; then
        print_status "OK" "CA certificate is valid"
    else
        print_status "FAIL" "CA certificate is invalid"
        return 1
    fi
    
    # Validate server certificate
    if openssl verify -CAfile "$ca_crt" "$server_crt" &>/dev/null; then
        print_status "OK" "Server certificate is valid and trusted by CA"
    else
        print_status "FAIL" "Server certificate validation failed"
        return 1
    fi
    
    # Validate client certificate
    if openssl verify -CAfile "$ca_crt" "$client_crt" &>/dev/null; then
        print_status "OK" "Client certificate is valid and trusted by CA"
    else
        print_status "FAIL" "Client certificate validation failed"
        return 1
    fi
    
    # Check certificate expiration
    local server_exp=$(openssl x509 -in "$server_crt" -noout -enddate | cut -d= -f2)
    local client_exp=$(openssl x509 -in "$client_crt" -noout -enddate | cut -d= -f2)
    
    print_status "OK" "Server certificate expires: $server_exp"
    print_status "OK" "Client certificate expires: $client_exp"
}

# Check Vault connectivity (optional)
check_vault_connectivity() {
    echo -e "\n${BLUE}Checking HashiCorp Vault Connectivity...${NC}"
    
    local vault_addr="${VAULT_ADDR:-}"
    local vault_token="${VAULT_TOKEN:-}"
    
    if [ -z "$vault_addr" ]; then
        print_status "WARN" "VAULT_ADDR not set - Vault integration disabled"
        return 0
    fi
    
    if [ -z "$vault_token" ]; then
        print_status "WARN" "VAULT_TOKEN not set - Vault integration disabled"
        return 0
    fi
    
    # Check if vault CLI is available
    if command -v vault &> /dev/null; then
        if vault status &>/dev/null; then
            print_status "OK" "Vault connectivity verified"
        else
            print_status "WARN" "Vault connectivity failed (server may be sealed or unreachable)"
        fi
    else
        print_status "WARN" "Vault CLI not available - cannot test connectivity"
    fi
}

# Test TLS configuration compilation
test_compilation() {
    echo -e "\n${BLUE}Testing TLS Configuration Compilation...${NC}"
    
    cd "$FOXHUNT_ROOT"
    
    # Test TLI compilation
    if cargo check -p tli --features tls &>/dev/null; then
        print_status "OK" "TLI compiles with TLS features"
    else
        print_status "FAIL" "TLI compilation failed with TLS features"
        return 1
    fi
    
    # Test trading service compilation
    if cargo check -p trading_service &>/dev/null; then
        print_status "OK" "Trading service compiles with TLS features"
    else
        print_status "FAIL" "Trading service compilation failed"
        return 1
    fi
}

# Run TLS integration tests
run_integration_tests() {
    echo -e "\n${BLUE}Running TLS Integration Tests...${NC}"
    
    cd "$FOXHUNT_ROOT"
    
    # Set test environment variables
    export TLS_CERT_PATH="$TLS_CERT_DIR/server.crt"
    export TLS_KEY_PATH="$TLS_CERT_DIR/server.key"
    export TLS_CA_CERT_PATH="$TLS_CERT_DIR/ca.crt"
    export REQUIRE_CLIENT_CERT="true"
    
    if cargo test tls_integration_tests --release &>/dev/null; then
        print_status "OK" "TLS integration tests passed"
    else
        print_status "WARN" "TLS integration tests failed (may be expected without running services)"
    fi
}

# Print configuration summary
print_configuration_summary() {
    echo -e "\n${BLUE}TLS Configuration Summary${NC}"
    echo "================================================"
    echo "Certificate Directory: $TLS_CERT_DIR"
    echo "Vault Cache Directory: $VAULT_CERT_CACHE"
    echo "Vault Address: ${VAULT_ADDR:-'Not set'}"
    echo "Environment Variables:"
    echo "  - TLS_CERT_PATH: ${TLS_CERT_PATH:-'Not set'}"
    echo "  - TLS_KEY_PATH: ${TLS_KEY_PATH:-'Not set'}"
    echo "  - TLS_CA_CERT_PATH: ${TLS_CA_CERT_PATH:-'Not set'}"
    echo "  - REQUIRE_CLIENT_CERT: ${REQUIRE_CLIENT_CERT:-'Not set'}"
    echo "  - USE_VAULT_TLS: ${USE_VAULT_TLS:-'Not set'}"
}

# Main execution
main() {
    check_permissions
    check_tonic_features
    check_certificate_directories
    generate_test_certificates
    validate_certificates
    check_vault_connectivity
    test_compilation
    run_integration_tests
    print_configuration_summary
    
    echo -e "\n${GREEN}🎉 TLS setup validation completed successfully!${NC}"
    echo -e "${YELLOW}Note: Test certificates generated are for development only.${NC}"
    echo -e "${YELLOW}Use proper certificates from your CA or Vault in production.${NC}"
}

# Run main function
main "$@"