#!/bin/bash
# Info: Polkadot Validator Key Management Script
# Author: Provisioning System

set -e

POLKADOT_BIN="{{ polkadot_validator.bin_path }}"
BASE_PATH="{{ polkadot_validator.base_path }}"
KEYSTORE_PATH="{{ polkadot_validator.keystore_path }}"
CONFIG_PATH="{{ polkadot_validator.config_path }}"
CHAIN="{{ polkadot_validator.network.chain }}"
RUN_USER="{{ polkadot_validator.run_user.name }}"

# Session keys file
SESSION_KEYS_FILE="{{ polkadot_validator.session_keys.keys_file | default('/var/lib/polkadot/session-keys') }}"
BACKUP_PATH="{{ polkadot_validator.security.backup_path | default('/var/backups/polkadot') }}"

echo "Polkadot Validator Key Management"
echo "================================="

# Function to generate session keys
generate_session_keys() {
    echo "Generating session keys..."
    
    # Call RPC to rotate keys
    RESULT=$(curl -H "Content-Type: application/json" -d '{"id":1, "jsonrpc":"2.0", "method": "author_rotateKeys", "params":[]}' http://localhost:9933 2>/dev/null | jq -r '.result' 2>/dev/null || echo "")
    
    if [ -n "$RESULT" ] && [ "$RESULT" != "null" ]; then
        echo "$RESULT" > "$SESSION_KEYS_FILE"
        echo "Session keys generated and saved to: $SESSION_KEYS_FILE"
        echo "Session keys: $RESULT"
        
        # Backup keys if enabled
        if [ "{{ polkadot_validator.security.backup_keys | lower }}" = "true" ]; then
            backup_keys
        fi
        
        return 0
    else
        echo "Failed to generate session keys via RPC. Is the node running?"
        return 1
    fi
}

# Function to backup keys
backup_keys() {
    echo "Backing up validator keys..."
    
    # Create backup directory
    mkdir -p "$BACKUP_PATH"
    BACKUP_DATE=$(date +%Y%m%d_%H%M%S)
    BACKUP_DIR="$BACKUP_PATH/keys_backup_$BACKUP_DATE"
    mkdir -p "$BACKUP_DIR"
    
    # Backup session keys
    if [ -f "$SESSION_KEYS_FILE" ]; then
        cp "$SESSION_KEYS_FILE" "$BACKUP_DIR/"
        echo "Session keys backed up"
    fi
    
    # Backup keystore (encrypted)
    if [ -d "$KEYSTORE_PATH" ]; then
        tar -czf "$BACKUP_DIR/keystore_backup.tar.gz" -C "$(dirname "$KEYSTORE_PATH")" "$(basename "$KEYSTORE_PATH")"
        echo "Keystore backed up"
    fi
    
    # Backup node key
    NODE_KEY_FILE="{{ polkadot_validator.network.node_key_file | default('/var/lib/polkadot/node-key') }}"
    if [ -f "$NODE_KEY_FILE" ]; then
        cp "$NODE_KEY_FILE" "$BACKUP_DIR/"
        echo "Node key backed up"
    fi
    
    # Set proper permissions
    chown -R "$RUN_USER:$RUN_USER" "$BACKUP_DIR"
    chmod -R 600 "$BACKUP_DIR"/*
    chmod 700 "$BACKUP_DIR"
    
    echo "Keys backed up to: $BACKUP_DIR"
}

# Function to restore keys from backup
restore_keys() {
    BACKUP_DIR="$1"
    
    if [ -z "$BACKUP_DIR" ] || [ ! -d "$BACKUP_DIR" ]; then
        echo "Usage: $0 restore <backup_directory>"
        echo "Available backups:"
        ls -la "$BACKUP_PATH"/keys_backup_* 2>/dev/null || echo "No backups found"
        return 1
    fi
    
    echo "Restoring keys from: $BACKUP_DIR"
    
    # Stop validator service for safety
    systemctl stop polkadot-validator 2>/dev/null || true
    
    # Restore session keys
    if [ -f "$BACKUP_DIR/session-keys" ]; then
        cp "$BACKUP_DIR/session-keys" "$SESSION_KEYS_FILE"
        echo "Session keys restored"
    fi
    
    # Restore keystore
    if [ -f "$BACKUP_DIR/keystore_backup.tar.gz" ]; then
        rm -rf "$KEYSTORE_PATH.old" 2>/dev/null || true
        mv "$KEYSTORE_PATH" "$KEYSTORE_PATH.old" 2>/dev/null || true
        tar -xzf "$BACKUP_DIR/keystore_backup.tar.gz" -C "$(dirname "$KEYSTORE_PATH")"
        echo "Keystore restored"
    fi
    
    # Restore node key
    NODE_KEY_FILE="{{ polkadot_validator.network.node_key_file | default('/var/lib/polkadot/node-key') }}"
    if [ -f "$BACKUP_DIR/node-key" ]; then
        cp "$BACKUP_DIR/node-key" "$NODE_KEY_FILE"
        echo "Node key restored"
    fi
    
    # Set proper permissions
    chown -R "$RUN_USER:$RUN_USER" "$KEYSTORE_PATH" "$SESSION_KEYS_FILE" "$NODE_KEY_FILE"
    chmod -R 600 "$KEYSTORE_PATH"/* "$SESSION_KEYS_FILE" "$NODE_KEY_FILE"
    
    echo "Keys restored successfully"
    echo "Starting validator service..."
    systemctl start polkadot-validator
}

# Function to verify session keys
verify_session_keys() {
    echo "Verifying session keys..."
    
    if [ ! -f "$SESSION_KEYS_FILE" ]; then
        echo "Session keys file not found: $SESSION_KEYS_FILE"
        return 1
    fi
    
    SESSION_KEYS=$(cat "$SESSION_KEYS_FILE")
    echo "Current session keys: $SESSION_KEYS"
    
    # Verify via RPC
    RESULT=$(curl -H "Content-Type: application/json" -d "{\"id\":1, \"jsonrpc\":\"2.0\", \"method\": \"author_hasSessionKeys\", \"params\":[\"$SESSION_KEYS\"]}" http://localhost:9933 2>/dev/null | jq -r '.result' 2>/dev/null || echo "false")
    
    if [ "$RESULT" = "true" ]; then
        echo "✅ Session keys are valid and loaded in the node"
    else
        echo "❌ Session keys are not loaded in the node"
        return 1
    fi
}

# Function to show current keys
show_keys() {
    echo "Current Validator Keys:"
    echo "======================"
    
    # Session keys
    if [ -f "$SESSION_KEYS_FILE" ]; then
        echo "Session keys: $(cat "$SESSION_KEYS_FILE")"
    else
        echo "Session keys: Not generated"
    fi
    
    # Node key (show public part only)
    NODE_KEY_FILE="{{ polkadot_validator.network.node_key_file | default('/var/lib/polkadot/node-key') }}"
    if [ -f "$NODE_KEY_FILE" ]; then
        if command -v "$POLKADOT_BIN" >/dev/null 2>&1; then
            PEER_ID=$("$POLKADOT_BIN" key inspect-node-key --file "$NODE_KEY_FILE" 2>/dev/null || echo "Unable to extract peer ID")
            echo "Node Peer ID: $PEER_ID"
        else
            echo "Node key: Present (run 'polkadot key inspect-node-key --file $NODE_KEY_FILE' to view peer ID)"
        fi
    else
        echo "Node key: Not generated"
    fi
    
    # Keystore info
    if [ -d "$KEYSTORE_PATH" ]; then
        KEY_COUNT=$(find "$KEYSTORE_PATH" -type f | wc -l)
        echo "Keystore keys: $KEY_COUNT files"
        echo "Keystore path: $KEYSTORE_PATH"
    else
        echo "Keystore: Not initialized"
    fi
}

# Function to set session keys on-chain
set_session_keys() {
    if [ ! -f "$SESSION_KEYS_FILE" ]; then
        echo "Session keys not found. Generate them first with: $0 generate"
        return 1
    fi
    
    SESSION_KEYS=$(cat "$SESSION_KEYS_FILE")
    echo "Setting session keys on-chain..."
    echo "Session keys: $SESSION_KEYS"
    echo ""
    echo "To set these keys on-chain:"
    echo "1. Go to https://polkadot.js.org/apps/#/staking/actions"
    echo "2. Click 'Set Session Key' for your stash account"
    echo "3. Paste the session keys: $SESSION_KEYS"
    echo "4. Submit the transaction"
    echo ""
    echo "Or use the Polkadot JS API:"
    echo "api.tx.session.setKeys('$SESSION_KEYS', '0x').signAndSend(account)"
}

# Function to rotate session keys
rotate_session_keys() {
    echo "Rotating session keys..."
    
    # Backup current keys
    if [ -f "$SESSION_KEYS_FILE" ]; then
        cp "$SESSION_KEYS_FILE" "$SESSION_KEYS_FILE.backup.$(date +%Y%m%d_%H%M%S)"
        echo "Current keys backed up"
    fi
    
    # Generate new keys
    generate_session_keys
    
    echo "Session keys rotated successfully"
    echo "Remember to update the keys on-chain!"
}

# Main command handling
case "${1:-help}" in
    "generate")
        generate_session_keys
        ;;
    "backup")
        backup_keys
        ;;
    "restore")
        restore_keys "$2"
        ;;
    "verify")
        verify_session_keys
        ;;
    "show")
        show_keys
        ;;
    "set")
        set_session_keys
        ;;
    "rotate")
        rotate_session_keys
        ;;
    "help"|*)
        echo "Usage: $0 [command]"
        echo ""
        echo "Commands:"
        echo "  generate     Generate new session keys"
        echo "  backup       Backup all validator keys"
        echo "  restore DIR  Restore keys from backup directory"
        echo "  verify       Verify current session keys"
        echo "  show         Show current keys information"
        echo "  set          Show instructions for setting keys on-chain"
        echo "  rotate       Rotate session keys (backup old, generate new)"
        echo "  help         Show this help message"
        echo ""
        echo "Examples:"
        echo "  $0 generate          # Generate new session keys"
        echo "  $0 verify            # Check if keys are loaded"
        echo "  $0 backup            # Backup all keys"
        echo "  $0 show              # Display key information"
        ;;
esac