#!/usr/bin/env bash

set -e

# Store validator node keys in Azure Key Vaults
# Each region's Key Vault stores the validator keys for that region

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
source "$SCRIPT_DIR/../lib/init.sh"
# Support both 48-node and 36-node configurations
NODES_FILE="/tmp/node_addresses_36.json"
if [ ! -f "$NODES_FILE" ]; then
    NODES_FILE="/tmp/node_addresses.json"
fi
KV_MAPPING="/tmp/key_vault_secrets_36.json"
if [ ! -f "$KV_MAPPING" ]; then
    KV_MAPPING="/tmp/key_vault_secrets.json"
fi

log_section "STORING VALIDATOR NODE KEYS IN AZURE KEY VAULTS"

# Dry-run mode (set DRY_RUN=1 to skip actual storage)
DRY_RUN="${DRY_RUN:-0}"

# Load node data using Python
python3 << PYTHON_EOF
import json
import subprocess
import sys
import os

# Support both 48-node and 36-node configurations
nodes_file = '/tmp/node_addresses_36.json'
if not os.path.exists(nodes_file):
    nodes_file = '/tmp/node_addresses.json'

kv_file = '/tmp/key_vault_secrets_36.json'
if not os.path.exists(kv_file):
    kv_file = '/tmp/key_vault_secrets.json'

with open(nodes_file, 'r') as f:
    nodes = json.load(f)

with open(kv_file, 'r') as f:
    kv_mapping = json.load(f)

print("╔════════════════════════════════════════════════════════════════╗")
print("║     STORING NODE KEYS IN AZURE KEY VAULTS                       ║")
print("╚════════════════════════════════════════════════════════════════╝")
print()

dry_run = os.environ.get('DRY_RUN', '0') == '1'
stored_count = 0
success_count = 0
failed_count = 0
skipped_count = 0
total_count = 0
errors = []

# Verify Key Vaults exist first
print("Verifying Key Vaults exist...")
kv_verified = {}
for region, config in kv_mapping.items():
    kv_name = config['key_vault_name']
    
    # Check if Key Vault exists
    result = subprocess.run(
        ['az', 'keyvault', 'show', '--name', kv_name, '--query', 'name', '-o', 'tsv'],
        capture_output=True,
        text=True
    )
    
    if result.returncode == 0 and result.stdout.strip() == kv_name:
        kv_verified[region] = True
        print(f"  ✓ {kv_name}")
    else:
        kv_verified[region] = False
        print(f"  ✗ {kv_name} (NOT FOUND)")
        errors.append(f"Key Vault not found: {kv_name} for region {region}")

print()

if not dry_run and len([k for k, v in kv_verified.items() if not v]) > 0:
    print("❌ ERROR: Some Key Vaults do not exist. Please deploy them first.")
    print("   Run: bash scripts/deployment/deploy-keyvaults-only.sh")
    sys.exit(1)

# Store secrets
for region, config in kv_mapping.items():
    kv_name = config['key_vault_name']
    region_nodes = config['nodes']
    
    if not kv_verified.get(region, False):
        print(f"⚠️  Skipping region {region} - Key Vault not found: {kv_name}")
        skipped_count += len(region_nodes)
        continue
    
    print(f"Region: {region}")
    print(f"  Key Vault: {kv_name}")
    
    for node in region_nodes:
        node_id = node['id']
        secret_name = f"validator-{node_id}-private-key"
        private_key = node.get('private_key_hex', '')
        
        if not private_key:
            print(f"    ⚠️  Skipping {secret_name} - no private key found")
            skipped_count += 1
            total_count += 1
            continue
        
        if dry_run:
            print(f"    [DRY RUN] Would store: {secret_name}")
        else:
            # Check if secret already exists
            check_result = subprocess.run(
                ['az', 'keyvault', 'secret', 'show', '--vault-name', kv_name, '--name', secret_name, '--query', 'name', '-o', 'tsv'],
                capture_output=True,
                text=True
            )
            
            if check_result.returncode == 0:
                print(f"    ⚠️  {secret_name} already exists - skipping")
                skipped_count += 1
            else:
                # Store private key as secret
                cmd = [
                    'az', 'keyvault', 'secret', 'set',
                    '--vault-name', kv_name,
                    '--name', secret_name,
                    '--value', private_key,
                    '--output', 'none'
                ]
                
                result = subprocess.run(cmd, capture_output=True, text=True)
                
                if result.returncode == 0:
                    print(f"    ✓ Stored: {secret_name}")
                    success_count += 1
                    stored_count += 1
                else:
                    error_msg = f"Failed to store {secret_name}: {result.stderr.strip()}"
                    print(f"    ✗ {error_msg}")
                    errors.append(error_msg)
                    failed_count += 1
        
        total_count += 1
    
    print()

print("=" * 64)
print("📊 SUMMARY")
print("=" * 64)
print()

print(f"Total nodes: {len(nodes)}")
print(f"Secrets to store: {total_count}")
if not dry_run:
    print(f"✓ Successfully stored: {success_count}")
    print(f"✗ Failed: {failed_count}")
    print(f"⚠️  Skipped: {skipped_count}")
    print(f"Key Vaults: {len(kv_mapping)}")
else:
    print(f"[DRY RUN] Would store: {total_count}")
print()

if errors:
    print("=" * 64)
    print("❌ ERRORS")
    print("=" * 64)
    for error in errors:
        print(f"  • {error}")
    print()

if failed_count > 0:
    sys.exit(1)

PYTHON_EOF

EXIT_CODE=$?

if [ $EXIT_CODE -eq 0 ]; then
    log_success "Key Vault storage complete!"
    echo ""
else
    log_error "Key Vault storage completed with errors"
    echo ""
    exit $EXIT_CODE
fi

