#!/bin/bash
# Info: Script to generate and manage Polkadot solochain keys
# Author: Provisioning System

set -e

POLKADOT_BIN_PATH="{{ polkadot_solochain.bin_path }}"
POLKADOT_NODE_BINARY="{{ polkadot_solochain.node_binary }}"
POLKADOT_BASE_PATH="{{ polkadot_solochain.base_path }}"
POLKADOT_CONFIG_PATH="{{ polkadot_solochain.config_path }}"
POLKADOT_RUN_USER="{{ polkadot_solochain.run_user.name }}"
CHAIN_SPEC_FILE="{{ polkadot_solochain.config_path }}/{{ polkadot_solochain.network.chain_id }}.json"

echo "Polkadot Solochain Key Management"
echo "================================="

# Function to generate Aura keys
generate_aura_key() {
    local seed="$1"
    local name="$2"
    
    echo "Generating Aura key for $name..."
    sudo -u "$POLKADOT_RUN_USER" "$POLKADOT_BIN_PATH/$POLKADOT_NODE_BINARY" key insert \
        --base-path "$POLKADOT_BASE_PATH" \
        --chain "$CHAIN_SPEC_FILE" \
        --scheme Sr25519 \
        --suri "$seed" \
        --key-type aura \
        --password-interactive < /dev/null
}

# Function to generate GRANDPA keys
generate_grandpa_key() {
    local seed="$1"
    local name="$2"
    
    echo "Generating GRANDPA key for $name..."
    sudo -u "$POLKADOT_RUN_USER" "$POLKADOT_BIN_PATH/$POLKADOT_NODE_BINARY" key insert \
        --base-path "$POLKADOT_BASE_PATH" \
        --chain "$CHAIN_SPEC_FILE" \
        --scheme Ed25519 \
        --suri "$seed" \
        --key-type gran \
        --password-interactive < /dev/null
}

# Function to generate session keys
generate_session_keys() {
    echo "Generating session keys..."
    
    # Generate random session keys
    AURA_SEED="$(openssl rand -hex 32)"
    GRANDPA_SEED="$(openssl rand -hex 32)"
    
    # Insert keys
    generate_aura_key "0x$AURA_SEED" "validator"
    generate_grandpa_key "0x$GRANDPA_SEED" "validator"
    
    # Save seeds for reference
    echo "AURA_SEED=0x$AURA_SEED" > "$POLKADOT_CONFIG_PATH/validator-seeds"
    echo "GRANDPA_SEED=0x$GRANDPA_SEED" >> "$POLKADOT_CONFIG_PATH/validator-seeds"
    chmod 600 "$POLKADOT_CONFIG_PATH/validator-seeds"
    chown "$POLKADOT_RUN_USER:$POLKADOT_RUN_USER" "$POLKADOT_CONFIG_PATH/validator-seeds"
    
    echo "Session keys generated and saved to $POLKADOT_CONFIG_PATH/validator-seeds"
}

# Function to generate development keys (Alice, Bob, etc.)
generate_dev_keys() {
    echo "Setting up development keys..."
    
    # Alice
    generate_aura_key "//Alice" "Alice"
    generate_grandpa_key "//Alice" "Alice"
    
    # Bob (if needed for multi-node setup)
    if [ "$1" = "multi" ]; then
        generate_aura_key "//Bob" "Bob"
        generate_grandpa_key "//Bob" "Bob"
        
        # Charlie
        generate_aura_key "//Charlie" "Charlie"
        generate_grandpa_key "//Charlie" "Charlie"
    fi
    
    echo "Development keys configured"
}

# Function to list existing keys
list_keys() {
    echo "Listing existing keys in keystore..."
    if [ -d "$POLKADOT_BASE_PATH/chains/{{ polkadot_solochain.network.chain_id }}/keystore" ]; then
        ls -la "$POLKADOT_BASE_PATH/chains/{{ polkadot_solochain.network.chain_id }}/keystore"
    else
        echo "No keystore found at $POLKADOT_BASE_PATH/chains/{{ polkadot_solochain.network.chain_id }}/keystore"
    fi
}

# Function to show public keys
show_public_keys() {
    echo "Extracting public keys..."
    if command -v jq >/dev/null 2>&1; then
        # Extract public keys from chain spec if available
        if [ -f "$CHAIN_SPEC_FILE" ]; then
            echo "Aura authorities:"
            jq -r '.genesis.runtime.aura.authorities[]?' "$CHAIN_SPEC_FILE" 2>/dev/null || echo "No Aura authorities found"
            
            echo "GRANDPA authorities:"
            jq -r '.genesis.runtime.grandpa.authorities[]?[0]' "$CHAIN_SPEC_FILE" 2>/dev/null || echo "No GRANDPA authorities found"
        fi
    else
        echo "jq not available - install jq to extract public keys from chain spec"
    fi
}

# Main command handling
case "${1:-help}" in
    "session")
        generate_session_keys
        ;;
    "dev")
        generate_dev_keys "${2:-single}"
        ;;
    "list")
        list_keys
        ;;
    "public")
        show_public_keys
        ;;
    "clean")
        echo "Removing all keys from keystore..."
        if [ -d "$POLKADOT_BASE_PATH/chains/{{ polkadot_solochain.network.chain_id }}/keystore" ]; then
            sudo -u "$POLKADOT_RUN_USER" rm -rf "$POLKADOT_BASE_PATH/chains/{{ polkadot_solochain.network.chain_id }}/keystore"/*
            echo "Keystore cleaned"
        else
            echo "No keystore found"
        fi
        ;;
    "help"|*)
        echo "Usage: $0 [command]"
        echo ""
        echo "Commands:"
        echo "  session     Generate random session keys for validator"
        echo "  dev [multi] Generate development keys (Alice, Bob, Charlie if multi)"
        echo "  list        List existing keys in keystore"
        echo "  public      Show public keys from chain specification"
        echo "  clean       Remove all keys from keystore"
        echo "  help        Show this help message"
        echo ""
        echo "Examples:"
        echo "  $0 dev          # Generate Alice keys for development"
        echo "  $0 dev multi    # Generate Alice, Bob, Charlie keys"
        echo "  $0 session      # Generate random validator keys"
        echo "  $0 list         # Show current keystore contents"
        ;;
esac