#!/bin/bash
# Info: Script to install Polkadot Validator
# Author: Provisioning System
# Release: 1.0
# Date: 2025-07-24

USAGE="install-polkadot-validator.sh"
[ "$1" == "-h" ] && echo "$USAGE" && exit 1

[ -r "env-polkadot-validator" ] && . ./env-polkadot-validator

POLKADOT_VERSION=${POLKADOT_VERSION:-latest}
POLKADOT_CHAIN=${POLKADOT_CHAIN:-polkadot}

# Determine architecture
ARCH="$(uname -m)"
case $ARCH in
    x86_64) ARCH="x86_64" ;;
    aarch64) ARCH="aarch64" ;;
    *) echo "Unsupported architecture: $ARCH" && exit 1 ;;
esac

# Set download URL based on version
if [ "$POLKADOT_VERSION" = "latest" ]; then
    POLKADOT_URL="https://github.com/paritytech/polkadot/releases/latest/download"
    POLKADOT_BINARY="polkadot"
else
    POLKADOT_URL="https://github.com/paritytech/polkadot/releases/download/${POLKADOT_VERSION}"
    POLKADOT_BINARY="polkadot"
fi

POLKADOT_BIN_PATH=${POLKADOT_BIN_PATH:-/usr/local/bin/polkadot}
POLKADOT_SYSTEMCTL_MODE=${POLKADOT_SYSTEMCTL_MODE:-enabled}

POLKADOT_CONFIG_PATH=${POLKADOT_CONFIG_PATH:-/etc/polkadot}
POLKADOT_WORK_PATH=${POLKADOT_WORK_PATH:-/var/lib/polkadot}
POLKADOT_BASE_PATH=${POLKADOT_BASE_PATH:-/var/lib/polkadot/data}
POLKADOT_KEYSTORE_PATH=${POLKADOT_KEYSTORE_PATH:-/var/lib/polkadot/keystore}

POLKADOT_RUN_USER=${POLKADOT_RUN_USER:-polkadot}
POLKADOT_RUN_GROUP=${POLKADOT_RUN_GROUP:-polkadot}
POLKADOT_RUN_USER_HOME=${POLKADOT_RUN_USER_HOME:-/home/polkadot}

POLKADOT_VALIDATOR_NAME=${POLKADOT_VALIDATOR_NAME:-polkadot-validator}

echo "Installing Polkadot Validator ${POLKADOT_VERSION}..."

# Check system requirements
echo "Checking system requirements..."

# Check CPU
CPU_CORES=$(nproc)
if [ "$CPU_CORES" -lt 8 ]; then
    echo "Warning: Polkadot validators require at least 8 CPU cores. Found: $CPU_CORES"
fi

# Check memory
TOTAL_MEM=$(free -g | awk '/^Mem:/{print $2}')
if [ "$TOTAL_MEM" -lt 32 ]; then
    echo "Warning: Polkadot validators require at least 32GB RAM. Found: ${TOTAL_MEM}GB"
fi

# Check storage
AVAILABLE_SPACE=$(df "$POLKADOT_BASE_PATH" 2>/dev/null | awk 'NR==2 {print $4}' || df / | awk 'NR==2 {print $4}')
REQUIRED_SPACE=2000000000  # 2TB in KB
if [ "$AVAILABLE_SPACE" -lt "$REQUIRED_SPACE" ]; then
    echo "Warning: Polkadot validators require at least 2TB NVMe SSD storage"
    echo "Available: $(($AVAILABLE_SPACE / 1024 / 1024))GB"
fi

# Install dependencies
echo "Installing dependencies..."
if command -v apt-get >/dev/null 2>&1; then
    apt-get update
    apt-get install -y curl ca-certificates jq ufw fail2ban unattended-upgrades prometheus-node-exporter
elif command -v yum >/dev/null 2>&1; then
    yum update -y
    yum install -y curl ca-certificates jq firewalld fail2ban dnf-automatic node_exporter
elif command -v dnf >/dev/null 2>&1; then
    dnf update -y
    dnf install -y curl ca-certificates jq firewalld fail2ban dnf-automatic golang-github-prometheus-node-exporter
else
    echo "Package manager not found. Please install dependencies manually."
    exit 1
fi

# Create user and group
if ! id "$POLKADOT_RUN_USER" &>/dev/null; then
    groupadd -r "$POLKADOT_RUN_GROUP"
    useradd -r -g "$POLKADOT_RUN_GROUP" -d "$POLKADOT_RUN_USER_HOME" -s /bin/bash -c "Polkadot validator user" "$POLKADOT_RUN_USER"
fi

# Create directories
mkdir -p "$POLKADOT_CONFIG_PATH"
mkdir -p "$POLKADOT_WORK_PATH"
mkdir -p "$POLKADOT_BASE_PATH"
mkdir -p "$POLKADOT_KEYSTORE_PATH"
mkdir -p "$POLKADOT_RUN_USER_HOME"

# Create backup directory if enabled
if [ "${POLKADOT_BACKUP_KEYS:-true}" = "true" ]; then
    BACKUP_PATH=${POLKADOT_BACKUP_PATH:-/var/backups/polkadot}
    mkdir -p "$BACKUP_PATH"
    chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$BACKUP_PATH"
    chmod 700 "$BACKUP_PATH"
fi

# Download and install Polkadot binary
cd /tmp
echo "Downloading Polkadot binary from ${POLKADOT_URL}/${POLKADOT_BINARY}..."
curl -L -o polkadot "${POLKADOT_URL}/${POLKADOT_BINARY}"

if [ ! -f "polkadot" ]; then
    echo "Failed to download Polkadot binary"
    exit 1
fi

# Install binary
chmod +x polkadot
mv polkadot "$(dirname "$POLKADOT_BIN_PATH")/"

# Generate node key if not exists
NODE_KEY_FILE="${POLKADOT_NODE_KEY_FILE:-$POLKADOT_WORK_PATH/node-key}"
if [ ! -f "$NODE_KEY_FILE" ]; then
    echo "Generating node key..."
    "$POLKADOT_BIN_PATH" key generate-node-key --file "$NODE_KEY_FILE"
fi

# Set ownership with strict permissions
chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$POLKADOT_WORK_PATH"
chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$POLKADOT_BASE_PATH"
chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$POLKADOT_KEYSTORE_PATH"
chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$POLKADOT_RUN_USER_HOME"
chown -R "$POLKADOT_RUN_USER:$POLKADOT_RUN_GROUP" "$POLKADOT_CONFIG_PATH"

# Set strict permissions for validator security
chmod 700 "$POLKADOT_WORK_PATH"
chmod 700 "$POLKADOT_KEYSTORE_PATH"
chmod 600 "$NODE_KEY_FILE"

# Build validator arguments
VALIDATOR_ARGS="--chain $POLKADOT_CHAIN"
VALIDATOR_ARGS="$VALIDATOR_ARGS --name $POLKADOT_VALIDATOR_NAME"
VALIDATOR_ARGS="$VALIDATOR_ARGS --base-path $POLKADOT_BASE_PATH"
VALIDATOR_ARGS="$VALIDATOR_ARGS --node-key-file $NODE_KEY_FILE"
VALIDATOR_ARGS="$VALIDATOR_ARGS --validator"

# Network configuration
VALIDATOR_ARGS="$VALIDATOR_ARGS --listen-addr ${POLKADOT_LISTEN_ADDR:-/ip4/0.0.0.0/tcp/30333}"

if [ -n "$POLKADOT_PUBLIC_ADDR" ]; then
    VALIDATOR_ARGS="$VALIDATOR_ARGS --public-addr $POLKADOT_PUBLIC_ADDR"
fi

if [ -n "$POLKADOT_BOOTNODES" ]; then
    IFS=',' read -ra BOOTNODES <<< "$POLKADOT_BOOTNODES"
    for bootnode in "${BOOTNODES[@]}"; do
        VALIDATOR_ARGS="$VALIDATOR_ARGS --bootnode $bootnode"
    done
fi

if [ -n "$POLKADOT_RESERVED_NODES" ]; then
    IFS=',' read -ra RESERVED <<< "$POLKADOT_RESERVED_NODES"
    for reserved in "${RESERVED[@]}"; do
        VALIDATOR_ARGS="$VALIDATOR_ARGS --reserved-node $reserved"
    done
fi

if [ "${POLKADOT_RESERVED_ONLY:-false}" = "true" ]; then
    VALIDATOR_ARGS="$VALIDATOR_ARGS --reserved-only"
fi

# RPC configuration (restricted for validator)
VALIDATOR_ARGS="$VALIDATOR_ARGS --rpc-bind-addr ${POLKADOT_RPC_BIND_ADDR:-127.0.0.1}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --rpc-port ${POLKADOT_RPC_PORT:-9944}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --rpc-methods ${POLKADOT_RPC_METHODS:-safe}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --rpc-max-connections ${POLKADOT_RPC_MAX_CONNECTIONS:-10}"

# Monitoring configuration
if [ "${POLKADOT_MONITORING_ENABLED:-true}" = "true" ]; then
    VALIDATOR_ARGS="$VALIDATOR_ARGS --prometheus-port ${POLKADOT_PROMETHEUS_PORT:-9615}"
    VALIDATOR_ARGS="$VALIDATOR_ARGS --prometheus-bind-addr ${POLKADOT_PROMETHEUS_BIND_ADDR:-127.0.0.1}"
fi

# Performance settings optimized for validator
VALIDATOR_ARGS="$VALIDATOR_ARGS --execution ${POLKADOT_EXECUTION:-wasm}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --wasm-execution ${POLKADOT_WASM_EXECUTION:-compiled}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --state-cache-size ${POLKADOT_STATE_CACHE_SIZE:-134217728}"
VALIDATOR_ARGS="$VALIDATOR_ARGS --db-cache ${POLKADOT_DB_CACHE:-2048}"

# Pruning (validators should keep more blocks)
VALIDATOR_ARGS="$VALIDATOR_ARGS --pruning ${POLKADOT_PRUNING:-1000}"
if [ "${POLKADOT_UNSAFE_PRUNING:-false}" = "true" ]; then
    VALIDATOR_ARGS="$VALIDATOR_ARGS --unsafe-pruning"
fi

# Telemetry
if [ "${POLKADOT_TELEMETRY_ENABLED:-true}" = "true" ]; then
    VALIDATOR_ARGS="$VALIDATOR_ARGS --telemetry-url '${POLKADOT_TELEMETRY_URL:-wss://telemetry.polkadot.io/submit/} ${POLKADOT_TELEMETRY_VERBOSITY:-0}'"
fi

# Logging
LOG_CONFIG="${POLKADOT_LOG_LEVEL:-info}"
if [ -n "$POLKADOT_LOG_TARGETS" ]; then
    LOG_CONFIG="$LOG_CONFIG,${POLKADOT_LOG_TARGETS}"
fi
VALIDATOR_ARGS="$VALIDATOR_ARGS --log $LOG_CONFIG"

# Create systemd service file
cat > /etc/systemd/system/polkadot-validator.service << EOF
[Unit]
Description=Polkadot Validator Node
Documentation=https://docs.polkadot.network/
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=$POLKADOT_RUN_USER
Group=$POLKADOT_RUN_GROUP
Environment=RUST_LOG=${POLKADOT_LOG_LEVEL:-info}
WorkingDirectory=$POLKADOT_WORK_PATH
ExecStart=$POLKADOT_BIN_PATH $VALIDATOR_ARGS
Restart=always
RestartSec=10

# Security settings (enhanced for validator)
NoNewPrivileges=true
PrivateTmp=true
ProtectSystem=strict
ProtectHome=true
ReadWritePaths=$POLKADOT_WORK_PATH $POLKADOT_BASE_PATH $POLKADOT_KEYSTORE_PATH $POLKADOT_CONFIG_PATH
CapabilityBoundingSet=CAP_NET_BIND_SERVICE
PrivateDevices=true
ProtectKernelTunables=true
ProtectKernelModules=true
ProtectControlGroups=true

# Resource limits (optimized for validator)
LimitNOFILE=65536
MemoryMax=16G

[Install]
WantedBy=multi-user.target
EOF

# Setup security if enabled
if [ "${POLKADOT_ENABLE_FIREWALL:-true}" = "true" ]; then
    echo "Setting up firewall..."
    
    if command -v ufw >/dev/null 2>&1; then
        # Ubuntu/Debian firewall
        ufw --force reset
        ufw default deny incoming
        ufw default allow outgoing
        
        # Allow SSH
        ufw allow ssh
        
        # Allow SSH from specific IPs if configured
        if [ -n "$POLKADOT_ALLOWED_SSH_IPS" ]; then
            ufw delete allow ssh
            IFS=',' read -ra SSH_IPS <<< "$POLKADOT_ALLOWED_SSH_IPS"
            for ip in "${SSH_IPS[@]}"; do
                ufw allow from "$ip" to any port 22
            done
        fi
        
        # Allow P2P port
        ufw allow 30333
        
        # Allow monitoring (localhost only)
        ufw allow from 127.0.0.1 to any port 9615
        ufw allow from 127.0.0.1 to any port 9944
        ufw allow from 127.0.0.1 to any port 9933
        
        ufw --force enable
        
    elif command -v firewall-cmd >/dev/null 2>&1; then
        # RHEL/CentOS firewall
        systemctl enable firewalld
        systemctl start firewalld
        
        firewall-cmd --permanent --add-port=30333/tcp
        firewall-cmd --permanent --add-service=ssh
        firewall-cmd --reload
    fi
fi

# Setup fail2ban if enabled
if [ "${POLKADOT_FAIL2BAN_ENABLED:-true}" = "true" ] && command -v fail2ban-client >/dev/null 2>&1; then
    echo "Configuring fail2ban..."
    systemctl enable fail2ban
    systemctl start fail2ban
fi

# Setup automatic updates if enabled
if [ "${POLKADOT_AUTO_UPDATES:-true}" = "true" ]; then
    echo "Enabling automatic security updates..."
    
    if command -v unattended-upgrades >/dev/null 2>&1; then
        # Ubuntu/Debian
        echo 'Unattended-Upgrade::Automatic-Reboot "false";' > /etc/apt/apt.conf.d/50unattended-upgrades-local
        systemctl enable unattended-upgrades
    elif command -v dnf >/dev/null 2>&1; then
        # RHEL/CentOS
        systemctl enable dnf-automatic.timer
        systemctl start dnf-automatic.timer
    fi
fi

# Install key management script
if [ -f "validator-keys.sh.j2" ]; then
    # This would be processed by template engine in real deployment
    cp validator-keys.sh.j2 "$POLKADOT_CONFIG_PATH/validator-keys.sh"
    chmod +x "$POLKADOT_CONFIG_PATH/validator-keys.sh"
    ln -sf "$POLKADOT_CONFIG_PATH/validator-keys.sh" "/usr/local/bin/polkadot-keys"
fi

# Enable and start service
systemctl daemon-reload
systemctl "$POLKADOT_SYSTEMCTL_MODE" polkadot-validator.service

if [ "$POLKADOT_SYSTEMCTL_MODE" = "enabled" ]; then
    systemctl start polkadot-validator.service
    
    # Wait for service to start
    sleep 10
fi

echo "=========================================="
echo "Polkadot Validator installation completed!"
echo "=========================================="
echo "Chain: $POLKADOT_CHAIN"
echo "Validator name: $POLKADOT_VALIDATOR_NAME"
echo "Service: polkadot-validator.service"
echo ""
echo "Node endpoints (localhost only for security):"
echo "WebSocket: ws://127.0.0.1:${POLKADOT_RPC_PORT:-9944}"
echo "HTTP RPC: http://127.0.0.1:${POLKADOT_HTTP_PORT:-9933}"
echo "Prometheus: http://127.0.0.1:${POLKADOT_PROMETHEUS_PORT:-9615}/metrics"
echo ""
echo "Configuration: $POLKADOT_CONFIG_PATH/"
echo "Data directory: $POLKADOT_BASE_PATH"
echo "Keystore: $POLKADOT_KEYSTORE_PATH"

# Show node peer ID
if [ -f "$NODE_KEY_FILE" ]; then
    PEER_ID=$("$POLKADOT_BIN_PATH" key inspect-node-key --file "$NODE_KEY_FILE" 2>/dev/null || echo "Unable to extract")
    echo "Node Peer ID: $PEER_ID"
fi

echo ""
echo "IMPORTANT: Next steps for validator setup:"
echo "1. Wait for node to sync completely"
echo "2. Generate session keys: polkadot-keys generate"
echo "3. Set up stash and controller accounts with sufficient DOT"
echo "4. Bond DOT tokens for staking"
echo "5. Set session keys on-chain: polkadot-keys set"
echo "6. Start validating from Polkadot.js Apps"
echo ""
echo "Security reminders:"
echo "- Keep your keystore and session keys backed up securely"
echo "- Monitor your validator for slashing risks"
echo "- Keep your node updated and online"
echo "- Never run duplicate validators with the same keys"

# Display service status
if systemctl is-active --quiet polkadot-validator.service; then
    echo "✅ Polkadot validator service is running"
    
    # Show initial sync status
    sleep 5
    echo ""
    echo "Checking initial sync status..."
    curl -s -H "Content-Type: application/json" -d '{"id":1, "jsonrpc":"2.0", "method": "system_health", "params":[]}' http://localhost:9933 | jq '.result' 2>/dev/null || echo "Node starting up..."
else
    echo "⚠️  Polkadot validator service status:"
    systemctl status polkadot-validator.service --no-pager -l
fi

# Cleanup
cd /
rm -rf /tmp/polkadot

echo ""
echo "Installation completed! Check the service status with:"
echo "systemctl status polkadot-validator"