{ pkgs ? import <nixpkgs> {} }:

let
    agenixSrc = fetchTarball {
        url = "https://github.com/ryantm/agenix/archive/main.tar.gz";
        sha256 = "103slb8xy5sb68zxjjbb9d0svq8xz751a7yrg6vrz5rh4374bzgl";
    };
in
pkgs.mkShell {
    buildInputs = [
        (pkgs.callPackage "${agenixSrc}/pkgs/agenix.nix" {})
    ];
    packages = with pkgs; [
        deploy-rs
        nano
        wireguard-tools
    ];

    EDITOR="nano";
    
    shellHook = ''
        rungcall() {
            echo "Running nix-collect-garbage -d on all nodes..."
            while read ip; do
                echo "============================================"
                echo "Running garbage collection on $ip"
                echo "============================================"
                ssh "$ip" "sudo nix-collect-garbage -d" < /dev/null 2>&1 | \
                while IFS= read -r line; do
                echo "[$ip] $line"
                done
                echo ""
            done < <(grep -o 'ip4 = "[0-9.]*/' nodes.nix | sed 's/ip4 = "//; s/\/.*//') 
        }
        runrebootall() {
            echo "This will reboot ALL nodes in the network!"
            echo "Nodes to be rebooted:"
            grep -o 'ip4 = "[0-9.]*/' nodes.nix | \
            sed 's/ip4 = "//; s/\/.*//' | \
            while read ip; do
                echo "  - $ip"
            done
            echo ""
            
            read -p "Are you sure you want to reboot all these nodes? (yes/no): " confirm
            
            if [ "$confirm" = "yes" ]; then
                echo "Rebooting all nodes..."
                while read ip; do
                echo "Rebooting $ip..."
                # Calling systemctl bypass molly-guard
                ssh "$ip" "sudo systemctl reboot" < /dev/null 2>&1 || echo "Yup, that failed"
                done < <(grep -o 'ip4 = "[0-9.]*/' nodes.nix | sed 's/ip4 = "//; s/\/.*//')
                echo "Reboot commands sent to all nodes."
            else
                echo "Reboot cancelled."
            fi
        }
        getallhk() {
            echo "Collecting SSH Ed25519 host keys from all nodes..."
            echo ""
            
            while read ip; do
                ssh "$ip" "cat /etc/ssh/ssh_host_ed25519_key.pub" < /dev/null 2>&1 || echo "Failed to get host key from $ip"
            done < <(grep -o 'ip4 = "[0-9.]*/' nodes.nix | sed 's/ip4 = "//; s/\/.*//') 
        }
        genwgkey() {
            echo "Generating WireGuard key pair..."
            umask 077
            local timestamp=$(date +%s%N)
            local private_key_file="/tmp/wg_private_$timestamp"
            local public_key_file="/tmp/wg_public_$timestamp"
            wg genkey > "$private_key_file"
            wg pubkey < "$private_key_file" > "$public_key_file"
            
            echo ""
            echo "============================================"
            echo "WireGuard Key Pair Generated:"
            echo "============================================"
            echo "Private Key:"
            cat "$private_key_file"
            echo ""
            echo "Public Key:"
            cat "$public_key_file"
            echo ""
            echo "============================================"
            
            shred -vfz -n 3 "$private_key_file" "$public_key_file" 2>/dev/null || {
                echo "Warning: shred not available, using rm..."
                rm -f "$private_key_file" "$public_key_file"
            }
            echo "Tpm Key files shreded."
        }
        deployparallel() {
            echo "Deploying to all nodes in parallel..."
            echo ""

            local deploy_args=()
            if [ $# -gt 0 ]; then
                deploy_args=("$@")
                echo "Deploy arguments: ${"\${deploy_args[*]}"}"
                echo ""
            fi
            
            local nodes=($(grep -E '^[[:space:]]*[a-zA-Z0-9_-]+-[a-zA-Z0-9_-]+-[a-zA-Z0-9_-]+-[a-zA-Z0-9_-]+[[:space:]]*=' nodes.nix | \
                           sed 's/^[[:space:]]*//; s/[[:space:]]*=.*//' | \
                           grep -v '^#'))
            
            if [ ${"\${#nodes[@]}"} -eq 0 ]; then
                echo "No nodes found in nodes.nix"
                return 1
            fi
            
            echo "Found ${"\${#nodes[@]}"} nodes:"
            for node in "${"\${nodes[@]}"}"; do
                echo " - $node"
            done
            echo ""
            
            read -p "Deploy to all these nodes? (yes/no): " confirm
            if [ "$confirm" != "yes" ]; then
                echo "Deployment cancelled."
                return 0
            fi
            
            echo "Starting parallel deployment..."
            echo "============================================"
            
            local tmpdir=$(mktemp -d)
            local pids=()

            for node in "${"\${nodes[@]}"}"; do
                {
                    echo "[$node] Starting deployment..."
                    deploy -s ".#$node" "${"\${deploy_args[@]}"}" 2>&1 | while IFS= read -r line; do
                        echo "[$node] $line"
                    done
                    echo "[$node] Deployment finished"
                } > "$tmpdir/$node.log" 2>&1 &
                pids+=($!)
            done
            
            while [ ${"\${#pids[@]}"} -gt 0 ]; do
                for i in "${"\${!pids[@]}"}"; do
                local pid=${"\${pids[i]}"}
                local node=${"\${nodes[i]}"}
                
                if ! kill -0 "$pid" 2>/dev/null; then
                    # Output is only displayed when the process end
                    if [ -f "$tmpdir/$node.log" ]; then
                        cat "$tmpdir/$node.log"
                    fi
                    unset 'pids[i]'
                    unset 'nodes[i]'
                fi
                done
                
                pids=(${"\${pids[@]}"})
                nodes=(${"\${nodes[@]}"})
                
                sleep 1
            done
            
            echo "============================================"
            echo "All deployments completed"
            
            rm -rf "$tmpdir"
        }

        
        export -f rungcall
        export -f runrebootall
        export -f getallhk
        export -f genwgkey
        export -f deployparallel
    
        echo "Welcome to Federez-LaSuite network deploy-rs shell environment!"
        echo "Available helper functions: rungcall, runrebootall, getallhk, genwgkey, deployparallel"
    '';
}