#!/usr/bin/env bash
# Setup JWT authentication for all ChainID 138 RPC containers (2503-2508)
# This script configures nginx with JWT authentication for each RPC container

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROXMOX_HOST="${PROXMOX_HOST:-192.168.11.10}"

# RPC containers requiring JWT authentication
declare -A RPC_CONTAINERS=(
    [2503]="besu-rpc-4-8a"
    [2504]="besu-rpc-4-1"
    [2505]="besu-rpc-luis-8a"
    [2506]="besu-rpc-luis-1"
    [2507]="besu-rpc-putu-8a"
    [2508]="besu-rpc-putu-1"
)

# IP addresses
declare -A RPC_IPS=(
    [2503]="192.168.11.253"
    [2504]="192.168.11.254"
    [2505]="192.168.11.255"
    [2506]="192.168.11.256"
    [2507]="192.168.11.257"
    [2508]="192.168.11.258"
)

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Check if container is running
check_container() {
    local vmid=$1
    local status=$(ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct status $vmid 2>/dev/null | awk '{print \$2}'" 2>/dev/null || echo "unknown")
    
    if [[ "$status" != "running" ]]; then
        return 1
    fi
    return 0
}

# Setup JWT authentication for a single container
setup_jwt_for_container() {
    local vmid=$1
    local container_name=$2
    local ip=$3
    
    log_info "Setting up JWT authentication for VMID $vmid ($container_name)..."
    
    # Check if container is running
    if ! check_container "$vmid"; then
        log_warn "Container $vmid is not running, skipping..."
        return 1
    fi
    
    # Install required packages
    log_info "  Installing nginx and dependencies..."
    ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct exec $vmid -- bash -c '
            export DEBIAN_FRONTEND=noninteractive
            apt-get update -qq
            apt-get install -y -qq nginx openssl python3 python3-pip || true
            pip3 install PyJWT cryptography 2>/dev/null || {
                apt-get install -y -qq python3-jwt python3-cryptography || true
            }
        '" || {
        log_warn "  Some packages may not be available, continuing..."
    }
    
    # Generate JWT secret
    log_info "  Generating JWT secret..."
    JWT_SECRET=$(ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct exec $vmid -- bash -c '
            if [ ! -f /etc/nginx/jwt_secret ]; then
                openssl rand -base64 32 > /etc/nginx/jwt_secret
                chmod 600 /etc/nginx/jwt_secret
            fi
            cat /etc/nginx/jwt_secret
        '")
    
    if [ -z "$JWT_SECRET" ]; then
        log_error "  Failed to generate JWT secret for $vmid"
        return 1
    fi
    
    log_success "  JWT secret generated: ${JWT_SECRET:0:20}..."
    
    # Create JWT validation script (using the same approach as configure-nginx-jwt-auth-simple.sh)
    log_info "  Creating JWT validation script..."
    ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct exec $vmid -- bash" <<'JWT_VALIDATOR_EOF'
cat > /usr/local/bin/jwt-validate.py <<'PYTHON_EOF'
#!/usr/bin/env python3
import sys
import os
import hmac
import hashlib
import base64
import json
import time

def base64url_decode(data):
    padding = 4 - len(data) % 4
    if padding != 4:
        data += '=' * padding
    return base64.urlsafe_b64decode(data)

def verify_jwt(token, secret):
    try:
        parts = token.split('.')
        if len(parts) != 3:
            return False, "Invalid token format"
        
        header_data = base64url_decode(parts[0])
        payload_data = base64url_decode(parts[1])
        
        header = json.loads(header_data)
        payload = json.loads(payload_data)
        
        if header.get('alg') != 'HS256':
            return False, "Unsupported algorithm"
        
        # Check expiry
        if 'exp' in payload:
            if int(payload['exp']) < time.time():
                return False, "Token expired"
        
        # Verify signature
        message = f"{parts[0]}.{parts[1]}"
        expected_signature = hmac.new(
            secret.encode('utf-8'),
            message.encode('utf-8'),
            hashlib.sha256
        ).digest()
        expected_signature_b64 = base64.urlsafe_b64encode(expected_signature).decode('utf-8').rstrip('=')
        
        if parts[2] != expected_signature_b64:
            return False, "Invalid signature"
        
        return True, payload
    except Exception as e:
        return False, str(e)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.exit(1)
    
    token = sys.argv[1]
    secret_file = '/etc/nginx/jwt_secret'
    
    if not os.path.exists(secret_file):
        sys.exit(1)
    
    with open(secret_file, 'r') as f:
        secret = f.read().strip()
    
    valid, result = verify_jwt(token, secret)
    if valid:
        print(json.dumps(result))
        sys.exit(0)
    else:
        sys.exit(1)
PYTHON_EOF
chmod +x /usr/local/bin/jwt-validate.py
JWT_VALIDATOR_EOF
    
    # Create JWT validation service
    log_info "  Creating JWT validation service..."
    ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct exec $vmid -- bash" <<'JWT_SERVICE_EOF'
cat > /etc/systemd/system/jwt-validator.service <<'SERVICE_EOF'
[Unit]
Description=JWT Validation Service
After=network.target

[Service]
Type=simple
ExecStart=/usr/bin/python3 /usr/local/bin/jwt-validate-service.py
Restart=always
RestartSec=5

[Install]
WantedBy=multi-user.target
SERVICE_EOF

cat > /usr/local/bin/jwt-validate-service.py <<'SERVICE_PYTHON_EOF'
#!/usr/bin/env python3
from http.server import HTTPServer, BaseHTTPRequestHandler
import json
import sys
import os

class JWTValidatorHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        if self.path == '/validate':
            content_length = int(self.headers.get('Content-Length', 0))
            post_data = self.rfc822_headers.get('Authorization', '')
            
            if post_data.startswith('Bearer '):
                token = post_data[7:]
                result = os.system(f'/usr/local/bin/jwt-validate.py "{token}" > /tmp/jwt_result.json 2>&1')
                
                if result == 0:
                    self.send_response(200)
                    self.end_headers()
                    self.wfile.write(b'OK')
                else:
                    self.send_response(401)
                    self.end_headers()
                    self.wfile.write(b'Unauthorized')
            else:
                self.send_response(401)
                self.end_headers()
                self.wfile.write(b'Unauthorized')
        else:
            self.send_response(404)
            self.end_headers()
    
    def log_message(self, format, *args):
        pass

if __name__ == '__main__':
    server = HTTPServer(('127.0.0.1', 8888), JWTValidatorHandler)
    server.serve_forever()
SERVICE_PYTHON_EOF
chmod +x /usr/local/bin/jwt-validate-service.py
JWT_SERVICE_EOF
    
    # Configure nginx
    log_info "  Configuring nginx..."
    ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no root@${PROXMOX_HOST} \
        "pct exec $vmid -- bash" <<NGINX_CONFIG_EOF
cat > /etc/nginx/sites-available/rpc-jwt <<'NGINX_EOF'
server {
    listen 443 ssl http2;
    listen [::]:443 ssl http2;
    server_name ${ip} ${container_name};

    ssl_certificate /etc/nginx/ssl/rpc.crt;
    ssl_certificate_key /etc/nginx/ssl/rpc.key;
    ssl_protocols TLSv1.2 TLSv1.3;
    ssl_ciphers HIGH:!aNULL:!MD5;
    ssl_prefer_server_ciphers on;

    add_header Strict-Transport-Security "max-age=31536000; includeSubDomains" always;
    add_header X-Frame-Options "SAMEORIGIN" always;
    add_header X-Content-Type-Options "nosniff" always;
    add_header X-XSS-Protection "1; mode=block" always;

    access_log /var/log/nginx/rpc-access.log;
    error_log /var/log/nginx/rpc-error.log;

    proxy_connect_timeout 300s;
    proxy_send_timeout 300s;
    proxy_read_timeout 300s;
    send_timeout 300s;

    # JWT authentication using auth_request
    location = /auth {
        internal;
        proxy_pass http://127.0.0.1:8888/validate;
        proxy_pass_request_body off;
        proxy_set_header Content-Length "";
        proxy_set_header X-Original-URI \$request_uri;
        proxy_set_header Authorization \$http_authorization;
    }

    # HTTP RPC endpoint
    location / {
        auth_request /auth;
        auth_request_set \$auth_status \$upstream_status;
        
        error_page 401 = @auth_failed;
        
        proxy_pass http://127.0.0.1:8545;
        proxy_http_version 1.1;
        proxy_set_header Host localhost;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
        proxy_set_header Connection "";
        proxy_buffering off;
        proxy_request_buffering off;
    }
    
    location @auth_failed {
        return 401 '{"jsonrpc":"2.0","error":{"code":-32000,"message":"Unauthorized. Missing or invalid JWT token. Use: Authorization: Bearer <token>"},"id":null}';
        add_header Content-Type application/json;
    }

    # Health check endpoint (no JWT required)
    location /health {
        access_log off;
        return 200 "healthy\n";
        add_header Content-Type text/plain;
    }
}
NGINX_EOF

# Create SSL directory and generate self-signed cert (or use existing)
mkdir -p /etc/nginx/ssl
if [ ! -f /etc/nginx/ssl/rpc.crt ]; then
    openssl req -x509 -nodes -days 365 -newkey rsa:2048 \
        -keyout /etc/nginx/ssl/rpc.key \
        -out /etc/nginx/ssl/rpc.crt \
        -subj "/CN=${container_name}/O=ChainID138/C=US"
fi

# Enable site
ln -sf /etc/nginx/sites-available/rpc-jwt /etc/nginx/sites-enabled/rpc-jwt
rm -f /etc/nginx/sites-enabled/default

# Test nginx configuration
nginx -t

# Start services
systemctl enable jwt-validator
systemctl start jwt-validator
systemctl restart nginx
NGINX_CONFIG_EOF
    
    log_success "  JWT authentication configured for VMID $vmid"
    
    # Save JWT secret for token generation
    echo "$JWT_SECRET" > "/tmp/jwt_secret_${vmid}.txt"
    log_info "  JWT secret saved to /tmp/jwt_secret_${vmid}.txt (for token generation)"
}

# Main execution
main() {
    log_info "Setting up JWT authentication for all ChainID 138 RPC containers..."
    echo ""
    
    local success_count=0
    local fail_count=0
    
    for vmid in "${!RPC_CONTAINERS[@]}"; do
        container_name="${RPC_CONTAINERS[$vmid]}"
        ip="${RPC_IPS[$vmid]}"
        
        if setup_jwt_for_container "$vmid" "$container_name" "$ip"; then
            ((success_count++))
        else
            ((fail_count++))
        fi
        echo ""
    done
    
    echo ""
    log_info "Summary:"
    log_success "  Successfully configured: $success_count containers"
    if [ $fail_count -gt 0 ]; then
        log_warn "  Failed/Skipped: $fail_count containers"
    fi
    
    echo ""
    log_info "Next steps:"
    log_info "  1. Generate JWT tokens using: ./scripts/generate-jwt-token-for-container.sh <VMID> <username> [expiry_days]"
    log_info "  2. Test JWT authentication on each container"
    log_info "  3. Distribute tokens to operators (Ali, Luis, Putu)"
}

main "$@"