#!/bin/bash
# Info: Automated session key rotation for Polkadot validator
# Author: Provisioning System

set -e

POLKADOT_BIN="{{ polkadot_validator.bin_path }}"
CONFIG_PATH="{{ polkadot_validator.config_path }}"
SESSION_KEYS_FILE="{{ polkadot_validator.session_keys.keys_file | default('/var/lib/polkadot/session-keys') }}"
ROTATION_INTERVAL="{{ polkadot_validator.session_keys.rotation_interval | default(86400) }}"
AUTO_ROTATE="{{ polkadot_validator.session_keys.auto_rotate | default(false) | lower }}"
RUN_USER="{{ polkadot_validator.run_user.name }}"

LOCK_FILE="/var/run/polkadot-session-rotation.lock"
LOG_FILE="/var/log/polkadot/session-rotation.log"

# Logging function
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$LOG_FILE"
}

# Check if rotation is needed
check_rotation_needed() {
    if [ ! -f "$SESSION_KEYS_FILE" ]; then
        log "No session keys found, rotation needed"
        return 0
    fi
    
    CURRENT_TIME=$(date +%s)
    FILE_TIME=$(stat -c %Y "$SESSION_KEYS_FILE" 2>/dev/null || echo "0")
    TIME_DIFF=$((CURRENT_TIME - FILE_TIME))
    
    if [ "$TIME_DIFF" -gt "$ROTATION_INTERVAL" ]; then
        log "Session keys are $TIME_DIFF seconds old, rotation needed (interval: $ROTATION_INTERVAL)"
        return 0
    else
        log "Session keys are $TIME_DIFF seconds old, no rotation needed"
        return 1
    fi
}

# Perform key rotation
rotate_keys() {
    log "Starting session key rotation..."
    
    # Create lock file
    if [ -f "$LOCK_FILE" ]; then
        log "Rotation already in progress (lock file exists)"
        return 1
    fi
    
    echo $$ > "$LOCK_FILE"
    trap 'rm -f "$LOCK_FILE"' EXIT
    
    # Backup current keys
    if [ -f "$SESSION_KEYS_FILE" ]; then
        BACKUP_FILE="$SESSION_KEYS_FILE.backup.$(date +%Y%m%d_%H%M%S)"
        cp "$SESSION_KEYS_FILE" "$BACKUP_FILE"
        log "Current keys backed up to: $BACKUP_FILE"
    fi
    
    # Generate new keys
    log "Generating new session keys..."
    RESULT=$(curl -s -H "Content-Type: application/json" -d '{"id":1, "jsonrpc":"2.0", "method": "author_rotateKeys", "params":[]}' http://localhost:9933 | jq -r '.result' 2>/dev/null || echo "")
    
    if [ -n "$RESULT" ] && [ "$RESULT" != "null" ]; then
        echo "$RESULT" > "$SESSION_KEYS_FILE"
        chown "$RUN_USER:$RUN_USER" "$SESSION_KEYS_FILE"
        chmod 600 "$SESSION_KEYS_FILE"
        log "New session keys generated: $RESULT"
        
        # Verify new keys
        VERIFY_RESULT=$(curl -s -H "Content-Type: application/json" -d "{\"id\":1, \"jsonrpc\":\"2.0\", \"method\": \"author_hasSessionKeys\", \"params\":[\"$RESULT\"]}" http://localhost:9933 | jq -r '.result' 2>/dev/null || echo "false")
        
        if [ "$VERIFY_RESULT" = "true" ]; then
            log "✅ New session keys verified successfully"
            
            # Send notification (if configured)
            send_notification "Session keys rotated successfully" "$RESULT"
            
            return 0
        else
            log "❌ Failed to verify new session keys"
            return 1
        fi
    else
        log "❌ Failed to generate new session keys"
        return 1
    fi
}

# Send notification
send_notification() {
    local message="$1"
    local keys="$2"
    
    # Log the notification
    log "NOTIFICATION: $message"
    
    # Send to syslog
    logger -t polkadot-validator "$message"
    
    # Additional notification methods can be added here
    # Examples: email, Slack, Discord, etc.
    
    # Example webhook notification (uncomment and configure)
    # if [ -n "$WEBHOOK_URL" ]; then
    #     curl -s -X POST "$WEBHOOK_URL" \
    #         -H "Content-Type: application/json" \
    #         -d "{\"text\":\"Polkadot Validator: $message\", \"keys\":\"$keys\"}" \
    #         || log "Failed to send webhook notification"
    # fi
}

# Check validator health
check_validator_health() {
    log "Checking validator health..."
    
    # Check if node is running
    if ! systemctl is-active --quiet polkadot-validator; then
        log "❌ Validator service is not running"
        return 1
    fi
    
    # Check node health via RPC
    HEALTH=$(curl -s -H "Content-Type: application/json" -d '{"id":1, "jsonrpc":"2.0", "method": "system_health", "params":[]}' http://localhost:9933 | jq -r '.result' 2>/dev/null)
    
    if [ -n "$HEALTH" ]; then
        IS_SYNCING=$(echo "$HEALTH" | jq -r '.isSyncing' 2>/dev/null || echo "true")
        PEERS=$(echo "$HEALTH" | jq -r '.peers' 2>/dev/null || echo "0")
        
        if [ "$IS_SYNCING" = "false" ] && [ "$PEERS" -gt 0 ]; then
            log "✅ Validator is healthy (synced, $PEERS peers)"
            return 0
        else
            log "⚠️ Validator may have issues (syncing: $IS_SYNCING, peers: $PEERS)"
            return 1
        fi
    else
        log "❌ Cannot check validator health (RPC not responding)"
        return 1
    fi
}

# Main execution
case "${1:-check}" in
    "rotate")
        log "Manual session key rotation requested"
        if check_validator_health; then
            rotate_keys
        else
            log "Skipping rotation due to validator health issues"
            exit 1
        fi
        ;;
    "check")
        if [ "$AUTO_ROTATE" = "true" ]; then
            log "Checking if automatic rotation is needed"
            if check_rotation_needed && check_validator_health; then
                rotate_keys
            fi
        else
            log "Automatic rotation is disabled"
        fi
        ;;
    "force")
        log "Forced session key rotation requested"
        rotate_keys
        ;;
    "health")
        check_validator_health
        ;;
    "status")
        log "Session key rotation status:"
        
        if [ -f "$SESSION_KEYS_FILE" ]; then
            CURRENT_TIME=$(date +%s)
            FILE_TIME=$(stat -c %Y "$SESSION_KEYS_FILE" 2>/dev/null || echo "0")
            TIME_DIFF=$((CURRENT_TIME - FILE_TIME))
            HOURS_OLD=$((TIME_DIFF / 3600))
            
            log "Current keys are $HOURS_OLD hours old"
            log "Rotation interval: $((ROTATION_INTERVAL / 3600)) hours"
            log "Auto rotation: $AUTO_ROTATE"
            
            if [ -f "$LOCK_FILE" ]; then
                log "Rotation in progress (PID: $(cat "$LOCK_FILE"))"
            else
                log "No rotation in progress"
            fi
        else
            log "No session keys found"
        fi
        ;;
    *)
        echo "Usage: $0 {check|rotate|force|health|status}"
        echo ""
        echo "Commands:"
        echo "  check    Check if rotation is needed and perform if auto-rotation enabled"
        echo "  rotate   Perform rotation if health checks pass"
        echo "  force    Force rotation regardless of timing"
        echo "  health   Check validator health"
        echo "  status   Show rotation status"
        echo ""
        echo "Configuration:"
        echo "  Auto rotation: $AUTO_ROTATE"
        echo "  Rotation interval: $((ROTATION_INTERVAL / 3600)) hours"
        echo "  Session keys file: $SESSION_KEYS_FILE"
        echo "  Log file: $LOG_FILE"
        exit 1
        ;;
esac