#!/bin/bash
#===============================================================================
# federation-distributed.sh - Distributed Training Across Federation
#
# Implements:
#   - Federated learning with gradient aggregation
#   - Parameter server architecture
#   - Byzantine fault tolerance (simplified)
#   - Consensus-based state synchronization
#===============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/lib/supervisor.sh" 2>/dev/null || true

#===============================================================================
# CONFIGURATION
#===============================================================================

FEDERATION_DIR="${FEDERATION_DIR:-/tmp/agi-federation}"
FEDERATION_STATE="${FEDERATION_DIR}/state"
FEDERATION_PARAMS="${FEDERATION_DIR}/params"
FEDERATION_GRADS="${FEDERATION_DIR}/gradients"

# Network configuration
PARAM_SERVER_PORT="${PARAM_SERVER_PORT:-9500}"
WORKER_BASE_PORT="${WORKER_BASE_PORT:-9510}"
DISCOVERY_PORT="${DISCOVERY_PORT:-9499}"

# Training configuration
LEARNING_RATE="${LEARNING_RATE:-0.001}"
MOMENTUM="${MOMENTUM:-0.9}"
GRADIENT_CLIP="${GRADIENT_CLIP:-1.0}"
SYNC_INTERVAL="${SYNC_INTERVAL:-10}"  # ticks between syncs

# Consensus configuration
QUORUM_SIZE="${QUORUM_SIZE:-2}"  # Minimum nodes for consensus
BYZANTINE_THRESHOLD="${BYZANTINE_THRESHOLD:-0.33}"  # Max fraction of faulty nodes

#===============================================================================
# INITIALIZATION
#===============================================================================

federation_init() {
    mkdir -p "$FEDERATION_DIR" "$FEDERATION_STATE" "$FEDERATION_PARAMS" "$FEDERATION_GRADS"

    # Initialize federation state
    cat > "${FEDERATION_STATE}/federation.json" <<EOF
{
    "federation_id": "$(uuidgen 2>/dev/null || echo $$-$(date +%s))",
    "created_at": "$(date -Iseconds)",
    "nodes": [],
    "leader": null,
    "epoch": 0,
    "global_tick": 0,
    "global_phi": 0.5,
    "status": "initializing"
}
EOF

    # Initialize global parameters (simplified: single vector)
    echo "0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5" > "${FEDERATION_PARAMS}/global.txt"

    echo "[FEDERATION] Initialized federation at $FEDERATION_DIR"
}

#===============================================================================
# NODE MANAGEMENT
#===============================================================================

declare -A NODES
declare -A NODE_STATES

register_node() {
    local node_id="$1"
    local node_host="${2:-localhost}"
    local node_port="${3:-$((WORKER_BASE_PORT + ${#NODES[@]}))}"

    NODES[$node_id]="${node_host}:${node_port}"
    NODE_STATES[$node_id]="active"

    # Update federation state
    if command -v jq &>/dev/null; then
        jq --arg id "$node_id" --arg addr "${node_host}:${node_port}" \
           '.nodes += [{"id": $id, "address": $addr, "status": "active"}]' \
           "${FEDERATION_STATE}/federation.json" > "${FEDERATION_STATE}/federation.json.tmp" \
           && mv "${FEDERATION_STATE}/federation.json.tmp" "${FEDERATION_STATE}/federation.json"
    fi

    echo "[FEDERATION] Registered node: $node_id at ${node_host}:${node_port}"
}

deregister_node() {
    local node_id="$1"

    unset "NODES[$node_id]"
    NODE_STATES[$node_id]="inactive"

    echo "[FEDERATION] Deregistered node: $node_id"
}

get_active_nodes() {
    local active=()
    for node_id in "${!NODE_STATES[@]}"; do
        if [[ "${NODE_STATES[$node_id]}" == "active" ]]; then
            active+=("$node_id")
        fi
    done
    echo "${active[@]}"
}

#===============================================================================
# PARAMETER SERVER
#===============================================================================

# Global parameter storage
declare -a GLOBAL_PARAMS
GLOBAL_EPOCH=0

load_global_params() {
    if [[ -f "${FEDERATION_PARAMS}/global.txt" ]]; then
        IFS=' ' read -r -a GLOBAL_PARAMS < "${FEDERATION_PARAMS}/global.txt"
    else
        GLOBAL_PARAMS=(0.5 0.5 0.5 0.5 0.5 0.5 0.5 0.5)
    fi
}

save_global_params() {
    echo "${GLOBAL_PARAMS[*]}" > "${FEDERATION_PARAMS}/global.txt"
}

# Apply gradient with momentum
declare -a MOMENTUM_BUFFER
apply_gradient() {
    local node_id="$1"
    local grad_file="${FEDERATION_GRADS}/${node_id}.txt"

    if [[ ! -f "$grad_file" ]]; then
        return
    fi

    local -a gradients
    IFS=' ' read -r -a gradients < "$grad_file"

    # Initialize momentum buffer if needed
    if [[ ${#MOMENTUM_BUFFER[@]} -eq 0 ]]; then
        for ((i = 0; i < ${#GLOBAL_PARAMS[@]}; i++)); do
            MOMENTUM_BUFFER[$i]=0
        done
    fi

    # Apply gradient with momentum
    for ((i = 0; i < ${#gradients[@]} && i < ${#GLOBAL_PARAMS[@]}; i++)); do
        # Gradient clipping
        local grad="${gradients[$i]}"
        if (( $(echo "$grad > $GRADIENT_CLIP" | bc -l) )); then
            grad="$GRADIENT_CLIP"
        elif (( $(echo "$grad < -$GRADIENT_CLIP" | bc -l) )); then
            grad="-$GRADIENT_CLIP"
        fi

        # Momentum update
        MOMENTUM_BUFFER[$i]=$(echo "scale=6; $MOMENTUM * ${MOMENTUM_BUFFER[$i]} + $grad" | bc)

        # Parameter update
        GLOBAL_PARAMS[$i]=$(echo "scale=6; ${GLOBAL_PARAMS[$i]} - $LEARNING_RATE * ${MOMENTUM_BUFFER[$i]}" | bc)
    done

    rm "$grad_file"
}

# Aggregate gradients from all nodes
aggregate_gradients() {
    local active_nodes=($(get_active_nodes))
    local num_nodes=${#active_nodes[@]}

    if [[ $num_nodes -lt $QUORUM_SIZE ]]; then
        echo "[FEDERATION] Not enough nodes for aggregation (${num_nodes}/${QUORUM_SIZE})"
        return 1
    fi

    # Collect gradients
    local -a sum_grads
    local collected=0

    for node_id in "${active_nodes[@]}"; do
        local grad_file="${FEDERATION_GRADS}/${node_id}.txt"
        if [[ -f "$grad_file" ]]; then
            local -a grads
            IFS=' ' read -r -a grads < "$grad_file"

            if [[ ${#sum_grads[@]} -eq 0 ]]; then
                sum_grads=("${grads[@]}")
            else
                for ((i = 0; i < ${#grads[@]}; i++)); do
                    sum_grads[$i]=$(echo "scale=6; ${sum_grads[$i]} + ${grads[$i]}" | bc)
                done
            fi
            ((collected++))
            rm "$grad_file"
        fi
    done

    if [[ $collected -eq 0 ]]; then
        return 1
    fi

    # Average gradients
    for ((i = 0; i < ${#sum_grads[@]}; i++)); do
        sum_grads[$i]=$(echo "scale=6; ${sum_grads[$i]} / $collected" | bc)
    done

    # Apply averaged gradient
    load_global_params
    for ((i = 0; i < ${#sum_grads[@]} && i < ${#GLOBAL_PARAMS[@]}; i++)); do
        GLOBAL_PARAMS[$i]=$(echo "scale=6; ${GLOBAL_PARAMS[$i]} - $LEARNING_RATE * ${sum_grads[$i]}" | bc)
    done
    save_global_params

    ((GLOBAL_EPOCH++))

    echo "[FEDERATION] Aggregated gradients from $collected nodes (epoch $GLOBAL_EPOCH)"
}

#===============================================================================
# WORKER NODE
#===============================================================================

worker_compute_gradient() {
    local node_id="$1"
    local input="${2:-}"

    load_global_params

    # Simulate forward pass and gradient computation
    local -a gradients
    for ((i = 0; i < ${#GLOBAL_PARAMS[@]}; i++)); do
        # Gradient = -(target - output) * input
        # Simplified: random gradient around 0
        local grad=$(echo "scale=6; ($(echo $RANDOM) / 32767 - 0.5) * 0.1" | bc)
        gradients+=("$grad")
    done

    # Save gradient
    echo "${gradients[*]}" > "${FEDERATION_GRADS}/${node_id}.txt"

    echo "[WORKER $node_id] Computed gradient"
}

worker_update_params() {
    local node_id="$1"

    # Fetch latest global params
    if [[ -f "${FEDERATION_PARAMS}/global.txt" ]]; then
        cp "${FEDERATION_PARAMS}/global.txt" "${FEDERATION_STATE}/${node_id}_params.txt"
        echo "[WORKER $node_id] Updated local params from global"
    fi
}

#===============================================================================
# CONSENSUS PROTOCOL
#===============================================================================

# Simplified PBFT-style consensus for state updates
propose_state() {
    local proposer="$1"
    local state_hash="$2"

    echo "[CONSENSUS] Node $proposer proposing state: $state_hash"

    # Broadcast proposal
    local active_nodes=($(get_active_nodes))
    local votes=0
    local total=${#active_nodes[@]}

    for node_id in "${active_nodes[@]}"; do
        # Simulate vote (in reality, would be network call)
        if [[ "$node_id" != "$proposer" ]]; then
            # 90% chance of agreement
            if [[ $((RANDOM % 100)) -lt 90 ]]; then
                ((votes++))
            fi
        fi
    done

    # Check quorum (2/3 + 1)
    local quorum=$(( (total * 2 / 3) + 1 ))
    if [[ $votes -ge $quorum ]]; then
        echo "[CONSENSUS] State accepted ($votes/$total votes, quorum=$quorum)"
        return 0
    else
        echo "[CONSENSUS] State rejected ($votes/$total votes, quorum=$quorum)"
        return 1
    fi
}

# Detect Byzantine nodes (simplified: check for outlier gradients)
detect_byzantine() {
    local active_nodes=($(get_active_nodes))
    local -a grad_norms

    for node_id in "${active_nodes[@]}"; do
        local grad_file="${FEDERATION_GRADS}/${node_id}.txt"
        if [[ -f "$grad_file" ]]; then
            local -a grads
            IFS=' ' read -r -a grads < "$grad_file"

            # Compute gradient norm
            local norm=0
            for g in "${grads[@]}"; do
                norm=$(echo "scale=6; $norm + $g * $g" | bc)
            done
            norm=$(echo "scale=6; sqrt($norm)" | bc -l 2>/dev/null || echo "0")

            grad_norms+=("$node_id:$norm")
        fi
    done

    # Find median norm
    local sorted_norms=($(for n in "${grad_norms[@]}"; do echo "${n#*:}"; done | sort -n))
    local median_idx=$(( ${#sorted_norms[@]} / 2 ))
    local median="${sorted_norms[$median_idx]:-1}"

    # Flag nodes with norm > 3x median
    for entry in "${grad_norms[@]}"; do
        local node_id="${entry%:*}"
        local norm="${entry#*:}"
        if (( $(echo "$norm > 3 * $median" | bc -l) )); then
            echo "[BYZANTINE] Detected potential Byzantine node: $node_id (norm=$norm, median=$median)"
            NODE_STATES[$node_id]="suspected"
        fi
    done
}

#===============================================================================
# TRAINING LOOP
#===============================================================================

run_training_round() {
    local round="$1"

    echo "=== Training Round $round ==="

    # Each worker computes gradient
    local active_nodes=($(get_active_nodes))
    for node_id in "${active_nodes[@]}"; do
        worker_compute_gradient "$node_id"
    done

    # Detect Byzantine nodes
    detect_byzantine

    # Aggregate gradients
    aggregate_gradients

    # Distribute updated params
    for node_id in "${active_nodes[@]}"; do
        worker_update_params "$node_id"
    done

    # Compute global phi
    load_global_params
    local sum=0
    for p in "${GLOBAL_PARAMS[@]}"; do
        sum=$(echo "scale=6; $sum + $p" | bc)
    done
    local mean=$(echo "scale=6; $sum / ${#GLOBAL_PARAMS[@]}" | bc)
    local phi=$(echo "scale=6; 1 / (1 + sqrt(($mean - 1)^2 + 0.01))" | bc -l 2>/dev/null || echo "0.5")

    echo "Global Phi: $phi (epoch $GLOBAL_EPOCH)"

    # Update federation state
    if command -v jq &>/dev/null; then
        jq --arg epoch "$GLOBAL_EPOCH" --arg phi "$phi" --arg tick "$round" \
           '.epoch = ($epoch | tonumber) | .global_phi = ($phi | tonumber) | .global_tick = ($tick | tonumber)' \
           "${FEDERATION_STATE}/federation.json" > "${FEDERATION_STATE}/federation.json.tmp" \
           && mv "${FEDERATION_STATE}/federation.json.tmp" "${FEDERATION_STATE}/federation.json"
    fi
}

run_training_daemon() {
    federation_init

    # Register some nodes
    register_node "node1"
    register_node "node2"
    register_node "node3"

    local round=0
    while true; do
        ((round++))
        run_training_round "$round"
        sleep "$SYNC_INTERVAL"
    done
}

#===============================================================================
# STATUS
#===============================================================================

federation_status() {
    echo "=== Federation Status ==="

    if [[ -f "${FEDERATION_STATE}/federation.json" ]]; then
        cat "${FEDERATION_STATE}/federation.json"
    fi

    echo ""
    echo "Nodes:"
    for node_id in "${!NODES[@]}"; do
        echo "  $node_id: ${NODES[$node_id]} (${NODE_STATES[$node_id]:-unknown})"
    done

    echo ""
    echo "Global Parameters:"
    if [[ -f "${FEDERATION_PARAMS}/global.txt" ]]; then
        cat "${FEDERATION_PARAMS}/global.txt"
    fi
}

#===============================================================================
# CLI
#===============================================================================

case "${1:-}" in
    init)
        federation_init
        ;;
    register)
        federation_init
        register_node "${2:-node1}" "${3:-localhost}" "${4:-9510}"
        ;;
    round)
        federation_init
        run_training_round "${2:-1}"
        ;;
    daemon)
        run_training_daemon
        ;;
    status)
        federation_status
        ;;
    *)
        echo "Usage: $0 {init|register|round|daemon|status}"
        ;;
esac